diff --git a/.github/workflows/check-coverage.yml b/.github/workflows/check-coverage.yml new file mode 100644 index 00000000..56bdf83d --- /dev/null +++ b/.github/workflows/check-coverage.yml @@ -0,0 +1,108 @@ +name: Check test coverage + +on: + push: + branches: + - main + - develop + pull_request: + branches: + - "*" + +jobs: + coverage: + name: Run test coverage check + runs-on: ubuntu-latest + container: precice/precice:nightly + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + path: micro-manager + + - name: Install dependencies + run: | + apt-get -qq update + apt-get -qq install python3-dev python3-venv git pkg-config + apt-get -qq install wget build-essential + + - name: Load Cache OpenMPI 5 + id: ompi-cache-load + uses: actions/cache/restore@v5 + with: + path: ~/openmpi + key: openmpi-v5-${{ runner.os }}-build + + - name: Build OpenMPI 5 + if: steps.ompi-cache-load.outputs.cache-hit != 'true' + run: | + wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.5.tar.gz + tar -xzf openmpi-5.0.5.tar.gz + cd openmpi-5.0.5 + mkdir -p ~/openmpi + ./configure --prefix=$HOME/openmpi + make -j$(nproc) + make install + + - name: Save OpenMPI 5 to cache + if: steps.ompi-cache-load.outputs.cache-hit != 'true' + uses: actions/cache/save@v5 + with: + path: ~/openmpi + key: openmpi-v5-${{ runner.os }}-build + + - name: Configure OpenMPI + run: | + cp -r ~/openmpi/* /usr/local/ + ldconfig + which mpiexec + mpiexec --version + + - name: Create a virtual environment and install Micro Manager with coverage + timeout-minutes: 6 + working-directory: micro-manager + run: | + python3 -m venv .venv + . .venv/bin/activate + pip install coverage + pip install .[sklearn,snapshot] + pip uninstall -y pyprecice + + - name: Run serial unit tests with coverage + working-directory: micro-manager/tests/unit + env: + PYTHONPATH: . + run: | + . ../../.venv/bin/activate + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_micro_manager + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_micro_simulation_crash_handling + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_adaptivity_serial + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_domain_decomposition + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_interpolation + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_hdf5_functionality + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_snapshot_computation + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_micro_simulation + + - name: Run tasking unit tests with coverage + working-directory: micro-manager/tests/unit + env: + PYTHONPATH: . + OMPI_ALLOW_RUN_AS_ROOT: "1" + OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: "1" + run: | + . ../../.venv/bin/activate + python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_tasking + + - name: Run parallel unit tests with coverage + working-directory: micro-manager/tests/unit + run: | + . ../../.venv/bin/activate + mpirun -n 2 --allow-run-as-root -x PYTHONPATH=. python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_adaptivity_parallel + mpirun -n 2 --allow-run-as-root -x PYTHONPATH=. python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_load_balancing + + - name: Combine coverage data + working-directory: micro-manager/tests/unit + run: | + . ../../.venv/bin/activate + python3 -m coverage combine + python3 -m coverage report --format=total | python3 ../../check_coverage.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a58ca012..5814bf93 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -17,7 +17,7 @@ jobs: - run: pip install --upgrade build - name: Build package run: pyproject-build - - uses: actions/upload-artifact@v6 + - uses: actions/upload-artifact@v7 with: name: dist path: dist @@ -33,7 +33,7 @@ jobs: id-token: write steps: - name: Download package - uses: actions/download-artifact@v7 + uses: actions/download-artifact@v8 with: name: dist path: dist diff --git a/.github/workflows/run-adaptivity-tests-parallel.yml b/.github/workflows/run-adaptivity-tests-parallel.yml index 83ee4c71..ad9d9407 100644 --- a/.github/workflows/run-adaptivity-tests-parallel.yml +++ b/.github/workflows/run-adaptivity-tests-parallel.yml @@ -59,14 +59,20 @@ jobs: with: path: micro-manager - - name: Install sudo for MPI - working-directory: micro-manager - run: | - apt-get -qq update - apt-get -qq install sudo + - name: Load Cache OpenMPI 5 + id: ompi-cache-load + uses: actions/cache/restore@v5 + with: + path: ~/openmpi + key: openmpi-v5-${{ runner.os }}-build + # If this fails, cache gets built in run-unit-tests - - name: Use mpi4py - uses: mpi4py/setup-mpi@v1 + - name: Configure OpenMPI + run: | + cp -r ~/openmpi/* /usr/local/ + ldconfig + which mpiexec + mpiexec --version - name: Install Dependencies working-directory: micro-manager @@ -94,12 +100,12 @@ jobs: run: | . .venv/bin/activate cd tests/unit - mpiexec -n 2 --allow-run-as-root python3 -m unittest test_global_adaptivity_lb.py + mpiexec -n 2 --allow-run-as-root python3 -m unittest test_load_balancing.py - - name: Run load balancing tests with 4 ranks + - name: Run load balancing unit tests with 4 ranks timeout-minutes: 3 working-directory: micro-manager run: | . .venv/bin/activate cd tests/unit - mpiexec -n 4 --allow-run-as-root --oversubscribe python3 -m unittest test_global_adaptivity_lb.py + mpiexec -n 4 --oversubscribe --allow-run-as-root python3 -m unittest test_load_balancing.py diff --git a/.github/workflows/run-unit-tests.yml b/.github/workflows/run-unit-tests.yml index 97fa7a55..a34eded2 100644 --- a/.github/workflows/run-unit-tests.yml +++ b/.github/workflows/run-unit-tests.yml @@ -20,6 +20,40 @@ jobs: run: | apt-get -qq update apt-get -qq install python3-dev python3-venv git pkg-config + apt-get -qq install wget build-essential + + - name: Load Cache OpenMPI 5 + id: ompi-cache-load + uses: actions/cache/restore@v5 + with: + path: ~/openmpi + key: openmpi-v5-${{ runner.os }}-build + + - name: Build OpenMPI 5 + if: steps.ompi-cache-load.outputs.cache-hit != 'true' + run: | + wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.5.tar.gz + tar -xzf openmpi-5.0.5.tar.gz + cd openmpi-5.0.5 + mkdir -p ~/openmpi + ./configure --prefix=$HOME/openmpi + make -j$(nproc) + make install + + - name: Save OpenMPI 5 to cache + if: steps.ompi-cache.outputs.cache-hit != 'true' + id: ompi-cache-store + uses: actions/cache/save@v5 + with: + path: ~/openmpi + key: openmpi-v5-${{ runner.os }}-build + + - name: Configure OpenMPI + run: | + cp -r ~/openmpi/* /usr/local/ + ldconfig + which mpiexec + mpiexec --version - name: Create a virtual environment and install Micro Manager in it timeout-minutes: 6 @@ -37,6 +71,16 @@ jobs: cd tests/unit python3 -m unittest test_micro_manager.py + - name: Install Micro Manager and run tasking unit test + working-directory: micro-manager + env: + OMPI_ALLOW_RUN_AS_ROOT: "1" + OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: "1" + run: | + . .venv/bin/activate + cd tests/unit + python3 -m unittest test_tasking.py + - name: Install Micro Manager and run interpolation unit test working-directory: micro-manager run: | diff --git a/.github/workflows/test-pip-install.yml b/.github/workflows/test-pip-install.yml new file mode 100644 index 00000000..a1db2fdb --- /dev/null +++ b/.github/workflows/test-pip-install.yml @@ -0,0 +1,51 @@ +name: Test pip installation and dependency check + +on: + push: + branches: + - develop + - main + pull_request: + branches: + - "*" + +jobs: + test-pip-install: + runs-on: ubuntu-latest + container: precice/precice:nightly + + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + path: micro-manager + + - name: Install dependencies + run: | + apt-get -qq update + apt-get -qq install python3-dev python3-venv git pkg-config + + - name: Create a virtual environment and install Micro Manager in it + working-directory: micro-manager + run: | + python3 -m venv .venv + . .venv/bin/activate + pip install . + + - name: Run dependency check with all dependencies present + working-directory: micro-manager + run: | + . .venv/bin/activate + micro-manager-precice --test-dependencies + + - name: Run dependency check with pyprecice uninstalled + working-directory: micro-manager + run: | + . .venv/bin/activate + pip uninstall -y pyprecice + if micro-manager-precice --test-dependencies; then + echo "ERROR: dependency check should have failed but passed" + exit 1 + else + echo "OK: dependency check correctly reported missing pyprecice" + fi diff --git a/.gitignore b/.gitignore index 4f2a6581..3a1b0f8a 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ dist precice-profiling/ precice-run/ *events.json +.coverage* diff --git a/CHANGELOG.md b/CHANGELOG.md index 74017e32..dc5540f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,22 @@ # Micro Manager changelog +## v0.9.0 + +- Refactored `DomainDecomposer` class and added a new variant of non-uniform decomposition [#243](https://github.com/precice/micro-manager/pull/243) +- Fixed load balancing configuration `partitioning` parameter setting [#245](https://github.com/precice/micro-manager/pull/245) +- Fixed comparison of zero values of type float32 and float64 in simulation deactivation [#244](https://github.com/precice/micro-manager/pull/244) +- Optimized norm calculations and further fixed lazy initialization [#241](https://github.com/precice/micro-manager/pull/241) +- Fixed lazy initialization for ranks without (active) micro simulations [#238](https://github.com/precice/micro-manager/pull/238) +- Added coverage testing and simulation interface tests [#225](https://github.com/precice/micro-manager/pull/225) +- Added `--test-dependencies` CLI flag to check if all required dependencies are correctly installed, with clear error messages listing missing packages and how to fix them [#221](https://github.com/precice/micro-manager/pull/221) +- Added load balancing based on micro simulation solve timings [#228](https://github.com/precice/micro-manager/pull/228) +- Fixed invalid value in division warning in L1rel/L2rel adaptivity when data contains zeros [#234](https://github.com/precice/micro-manager/pull/234) +- Fixed duplicate micro simulations for macro-points on rank boundaries by filtering coordinates already claimed by lower-ranked ranks [#230](https://github.com/precice/micro-manager/pull/230) +- Exposed `MicroSimulationInterface` as a public abstract base class for user subclassing [#224](https://github.com/precice/micro-manager/pull/224) +- Added option to use compute instances to reduce memory consumption [#226](https://github.com/precice/micro-manager/pull/226) +- Added support to run micro simulations in separate processes with workers [#219](https://github.com/precice/micro-manager/pull/219) +- Added abstraction layers to micro simulations to support more features [#218](https://github.com/precice/micro-manager/pull/218) + ## v0.8.0 - Conformed to naming standard in precice/tutorials [#215](https://github.com/precice/micro-manager/pull/215) diff --git a/check_coverage.py b/check_coverage.py new file mode 100644 index 00000000..969b6dfd --- /dev/null +++ b/check_coverage.py @@ -0,0 +1,32 @@ +""" +Script to check if test coverage is above a predefined threshold. +Reads total coverage percentage from stdin (output of coverage report --format=total). +""" + +import sys + +THRESHOLD = 60 # Minimum required coverage percentage + + +if __name__ == "__main__": + try: + coverage = int(sys.stdin.read().strip()) + except ValueError: + print("Error: could not parse coverage value from stdin.") + sys.exit(1) + + print("Total coverage: {}%".format(coverage)) + if coverage < THRESHOLD: + print( + "Coverage {}% is below the required threshold of {}%.".format( + coverage, THRESHOLD + ) + ) + sys.exit(1) + else: + print( + "Coverage {}% meets the required threshold of {}%.".format( + coverage, THRESHOLD + ) + ) + sys.exit(0) diff --git a/docs/configuration.md b/docs/configuration.md index 123d6d23..a7940cea 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -11,12 +11,13 @@ The Micro Manager is configured with a JSON file. Several parameters can be set. ## Micro Manager Configuration -Parameter | Description | Default ---- | --- -`micro_file_name` | Path to the file containing the Python importable micro simulation class. If the file is not in the working directory, give the relative path from the directory where the Micro Manager is executed. | - -`output_directory` | Path to output directory for logging and performance metrics. Directory is created if not existing already. | `.` -`memory_usage_output_type` | Set to either `local`, `global`, or `all`. `local` outputs rank-wise peak memory usage. `global` outputs global averaged peak memory usage. `all` outputs both local and global levels. | Empty string. -`memory_usage_output_n` | Interval of output. | 1 +| Parameter | Description | Default | +|----------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------| +| `micro_file_name` | Path to the file containing the Python importable micro simulation class. If the file is not in the working directory, give the relative path from the directory where the Micro Manager is executed. | - | +| `micro_stateless` | Boolean if micro simulation is stateless allowing model instancing. | False | +| `output_directory` | Path to output directory for logging and performance metrics. Directory is created if not existing already. | `.` | +| `memory_usage_output_type` | Set to either `local`, `global`, or `all`. `local` outputs rank-wise peak memory usage. `global` outputs global averaged peak memory usage. `all` outputs both local and global levels. | Empty string. | +| `memory_usage_output_n` | Interval of output. | 1 | All output is to a CSV file with the peak memory usage (RSS) in every time window, in MBs. @@ -24,31 +25,35 @@ Apart from the base settings, there are three main sections in the configuration ## Coupling Parameters -Parameter | Description ---- | --- -`precice_config_file_name` | Path to the preCICE XML configuration file from the current working directory. -`macro_mesh_name` | Name of the macro mesh as stated in the preCICE configuration. -`read_data_names` | List with the names of the data to be read from preCICE. -`write_data_names` | List with the names of the data to be written to preCICE. +| Parameter | Description | +|----------------------------|--------------------------------------------------------------------------------| +| `precice_config_file_name` | Path to the preCICE XML configuration file from the current working directory. | +| `macro_mesh_name` | Name of the macro mesh as stated in the preCICE configuration. | +| `read_data_names` | List with the names of the data to be read from preCICE. | +| `write_data_names` | List with the names of the data to be written to preCICE. | ## Simulation Parameters -Parameter | Description | Default ---- | --- | --- -`macro_domain_bounds` | Minimum and maximum bounds of the macro-domain, having the format `[xmin, xmax, ymin, ymax, zmin, zmax]` in 3D and `[xmin, xmax, ymin, ymax]` in 2D. | - -`decomposition` | List of number of ranks in each axis with format `[xranks, yranks, zranks]` in 3D and `[xranks, yranks]` in 2D. | `[1, 1, 1]` or `[1, 1]` -`micro_dt` | Initial time window size (dt) of the micro simulation. | - -`adaptivity` | Set `true` for simulations with adaptivity. See section on [adaptivity](#adaptivity). | `false` -`load_balancing` | Set `true` for load balancing. See section on [load balancing](#load-balancing). | `false` +| Parameter | Description | Default | +| --- | --- | --- | +| `macro_domain_bounds` | Minimum and maximum bounds of the macro-domain, having the format `[xmin, xmax, ymin, ymax, zmin, zmax]` in 3D and `[xmin, xmax, ymin, ymax]` in 2D. | - | +| `decomposition` | List of number of ranks in each axis with format `[xranks, yranks, zranks]` in 3D and `[xranks, yranks]` in 2D. | `[1, 1, 1]` or `[1, 1]` | +| `decomposition_type` | Type of domain decomposition. Either `uniform` or `nonuniform`. | `uniform` | +| `minimum_access_region_size` | If `nonuniform` decomposition, optionally set a minimum domain width in each axis. Format `[xmin, ymin, zmin]` | - | +| `micro_dt` | Initial time window size (dt) of the micro simulation. | - | +| `adaptivity` | Set `true` for simulations with adaptivity. See section on [adaptivity](#adaptivity). | `false` | +| `load_balancing` | Set `true` for load balancing. See section on [load balancing](#load-balancing). | `false` | The total number of partitions ranks in the `decomposition` list should be the same as the number of ranks in the `mpirun` or `mpiexec` command. +Non-uniform domain decomposition is based on a geometric progression. + ## Diagnostics -Parameter | Description | Default ---- | --- | --- -`data_from_micro_sims` | Dictionary with the names of the data from the micro simulation to be written to VTK files as keys and `"scalar"` or `"vector"` as values. | - -`micro_output_n` | Frequency of calling the optional output functionality of the micro simulation in terms of number of time steps. | 1 +| Parameter | Description | Default | +|------------------------|--------------------------------------------------------------------------------------------------------------------------------------------|---------| +| `data_from_micro_sims` | Dictionary with the names of the data from the micro simulation to be written to VTK files as keys and `"scalar"` or `"vector"` as values. | - | +| `micro_output_n` | Frequency of calling the optional output functionality of the micro simulation in terms of number of time steps. | 1 | ### Adding diagnostics in the preCICE XML configuration @@ -70,20 +75,20 @@ See the [adaptivity](tooling-micro-manager-adaptivity.html) documentation for a To turn on adaptivity, set `"adaptivity": true` in `simulation_params`. Then under `adaptivity_settings` set the following variables: -Parameter | Description | Default ---- | --- | --- -`type` | Set to either `local` or `global`. The type of adaptivity matters when the Micro Manager is run in parallel. `local` means comparing micro simulations within a local partitioned domain for similarity. `global` means comparing micro simulations from all partitions, so over the entire domain. | None -`data` | List of names of data which are to be used to calculate if micro-simulations are similar or not. For example `["temperature", "porosity"]`. | - -`adaptivity_every_n_time_windows` | Interval of adaptivity computation. | 1 -`output_type` | Set to either `local`, `global`, or `all`. `local` outputs rank-wise adaptivity metrics. `global` outputs global averaged metrics. `all` outputs both local and global metrics. | Empty string. -`output_n` | Frequency of output of adaptivity metrics. | 1 -`history_param` | History parameter $$ \Lambda $$, set as $$ \Lambda >= 0 $$. | 0.5 -`coarsening_constant` | Coarsening constant $$ C_c $$, set as $$ 0 =< C_c < 1 $$. | 0.5 -`refining_constant` | Refining constant $$ C_r $$, set as $$ 0 =< C_r < 1 $$. | 0.5 -`every_implicit_iteration` | If `true`, adaptivity is calculated in every implicit iteration.
If False, adaptivity is calculated once at the start of the time window and then reused in every implicit time iteration. | `false` -`similarity_measure` | Similarity measure to be used for adaptivity. Can be either `L1`, `L2`, `L1rel` or `L2rel`. By default, `L1` is used. The `rel` variants calculate the respective relative norms. This parameter is *optional*. | `L2rel` -`lazy_initialization` | Set to `true` to lazily create and initialize micro simulations. If selected, micro simulation objects are created only when the micro simulation is activated for the first time. | `false` -`load_balancing` | Set to `true` to dynamically balance simulations for parallel runs. See [load balancing settings](#load-balancing) below. | `false` +| Parameter | Description | Default | +|-----------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------| +| `type` | Set to either `local` or `global`. The type of adaptivity matters when the Micro Manager is run in parallel. `local` means comparing micro simulations within a local partitioned domain for similarity. `global` means comparing micro simulations from all partitions, so over the entire domain. | None | +| `data` | List of names of data which are to be used to calculate if micro-simulations are similar or not. For example `["temperature", "porosity"]`. | - | +| `adaptivity_every_n_time_windows` | Interval of adaptivity computation. | 1 | +| `output_type` | Set to either `local`, `global`, or `all`. `local` outputs rank-wise adaptivity metrics. `global` outputs global averaged metrics. `all` outputs both local and global metrics. | Empty string. | +| `output_n` | Frequency of output of adaptivity metrics. | 1 | +| `history_param` | History parameter $$ \Lambda $$, set as $$ \Lambda >= 0 $$. | 0.5 | +| `coarsening_constant` | Coarsening constant $$ C_c $$, set as $$ 0 =< C_c < 1 $$. | 0.5 | +| `refining_constant` | Refining constant $$ C_r $$, set as $$ 0 =< C_r < 1 $$. | 0.5 | +| `every_implicit_iteration` | If `true`, adaptivity is calculated in every implicit iteration.
If False, adaptivity is calculated once at the start of the time window and then reused in every implicit time iteration. | `false` | +| `similarity_measure` | Similarity measure to be used for adaptivity. Can be either `L1`, `L2`, `L1rel` or `L2rel`. By default, `L1` is used. The `rel` variants calculate the respective relative norms. This parameter is *optional*. | `L2rel` | +| `lazy_initialization` | Set to `true` to lazily create and initialize micro simulations. If selected, micro simulation objects are created only when the micro simulation is activated for the first time. | `false` | +| `load_balancing` | Set to `true` to dynamically balance simulations for parallel runs. See [load balancing settings](#load-balancing) below. | `false` | Example of adaptivity configuration is @@ -110,10 +115,11 @@ See the [model adaptivity](tooling-micro-manager-model-adaptivity.html) document To turn on model adaptivity, set `"model_adaptivity": true` in `simulation_params`. Then under `model_adaptivity_settings` set the following variables: -Parameter | Description ---- | --- -`micro_file_names` | List of paths to the files containing the Python importable micro simulation classes. If the files are not in the working directory, give the relative path from the directory where the Micro Manager is executed. Requires a minimum of 2 files. -`switching_function` | Path to the file containing the Python importable switching function. If the file is not in the working directory, give the relative path from the directory where the Micro Manager is executed. +| Parameter | Description | +|----------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `micro_file_names` | List of paths to the files containing the Python importable micro simulation classes. If the files are not in the working directory, give the relative path from the directory where the Micro Manager is executed. Requires a minimum of 2 files. | +| `switching_function` | Path to the file containing the Python importable switching function. If the file is not in the working directory, give the relative path from the directory where the Micro Manager is executed. | +| `micro_stateless` | List of boolean values, whether the respective micro simulation model is stateless and can use model instancing. | Example of model adaptivity configuration is @@ -122,7 +128,8 @@ Example of model adaptivity configuration is "model_adaptivity": true, "model_adaptivity_settings": { "micro_file_names": ["python-dummy/micro_dummy", "python-dummy/micro_dummy", "python-dummy/micro_dummy"], - "switching_function": "mada_switcher" + "switching_function": "mada_switcher", + "micro_stateless": [False, True, True] } } ``` @@ -153,13 +160,34 @@ The Micro Manager uses the output functionality of preCICE, hence these data set ## Load balancing -Under `load_balancing_settings`, the following parameters can be set +Load balancing can be activated by setting `load_balancing` to true. +It balances based on either the elapsed time required to solve the prior iteration `type="time""` or the number of active simulations `type=active`. +One Initial load balancing step is performed, prior to any computation (assuming equal workload for time based load balancing or the current active counts for `active` load balancing.). +Subsequently, in the following iteration another load balancing step is performed based. (This is mainly for the time based balancing to use the just acquired timings.) +Afterwards balancing is performed `every_n_time_windows`. +Upon activation, further configuration must be provided in `load_balancing_settings`. +The following parameters can be set + +| Parameter | Description | Default | +|-------------------------|--------------------------------------------------|----------| +| `every_n_time_windows` | Frequency of balancing the simulations. | `1` | +| `partitioning` | Partitioning Algorithm. Options: ["lpt"]. | `"lpt"` | +| `type` | Load balancing type. Options: ["time", "active"] | `"time"` | +| `threshold` | Threshold parameter | `0` | +| `balance_inactive_sims` | Balance inactive simulations | `False` | -Parameter | Description | Default ---- | --- | --- -`every_n_time_windows` | Frequency of balancing the simulations. | `1` -`balancing_threshold` | Integer threshold value. | `0` -`balance_inactive_sims` | If `true`, inactive simulations associated to redistributed active simulations are moved to the new ranks of the active simulations. See [balance inactive simulations](tooling-micro-manager-adaptivity.html#balance-inactive-simulations). | `false` +```json +"simulation_params": { + "load_balancing": true, + "load_balancing_settings": { + "every_n_time_windows": 5, + "partitioning": "lpt", + "type": "active", + "threshold": 2, + "balance_inactive_sims": true + } +} +``` ## Interpolate a crashed micro simulation diff --git a/docs/model-adaptivity.md b/docs/model-adaptivity.md index 3102dfaa..9596ccdc 100644 --- a/docs/model-adaptivity.md +++ b/docs/model-adaptivity.md @@ -115,6 +115,11 @@ class MicroSimulation: # Name is fixed It will be called with frequency set by configuration option `simulation_params: micro_output_n` This function is *optional*. """ + + def get_global_id(self): + """ + Return the assigned global id. + """ ``` For this the default MicroSimulation still serves as the model interface, while the `(set)|(get)_state()` methods @@ -124,40 +129,43 @@ is likely to be the full order model, while subsequent ones are ROMs. ```python def switching_function( - resolutions: np.ndarray, - locations: np.ndarray, + resolution: int, + location: np.ndarray, t: float, - inputs: list[dict], + input: dict, prev_output: dict, - active: np.ndarray, -) -> np.ndarray: +) -> int: """ Switching interface function, use as reference Parameters ---------- - resolutions : np.array - shape(N,) - Array with resolution information as get_sim_class_resolution would return for a sim obj. - locations : np.array - shape(N,D) - Array with gaussian points for all sims. D is the mesh dimension. + resolution : int + Current resolution as get_sim_class_resolution would return for the respective sim obj. + location : np.array - shape(D,) + Array with gaussian points. t : float Current time in simulation. - inputs : list[dict] - List of input objects. - prev_output : [None, dict-like] + input : dict + Input object. + prev_output : [None, dict] Contains the outputs of the previous model evaluation. - active : np.array - shape(N,) - Bool array indicating whether the model is active or not. + Returns + ------- + switch_direction: int + 0 if resolution should not change + -1 if resolution should increase + 1 if resolution should decrease """ - return np.zeros_like(resolutions) + return 0 ``` -The switching of models is governed by the `switching_function`. -The output is expected to be a np.ndarray of shape (N,) and is interpreted in the following manner: +The switching of models is governed by the `switching_function`, which is evaluated for each micro simulation. +The output is expected to be an integer and is interpreted in the following manner: -Value | Action ---- | --- -0 | No resolution change --1 | Increase model fidelity by one (go back one in list) -1 | Decrease model fidelity by one (go one ahead in list) +| Value | Action | +|-------|-------------------------------------------------------| +| 0 | No resolution change | +| -1 | Increase model fidelity by one (go back one in list) | +| 1 | Decrease model fidelity by one (go one ahead in list) | diff --git a/examples/cpp-dummy/micro_cpp_dummy.cpp b/examples/cpp-dummy/micro_cpp_dummy.cpp index f4aaad6d..24788483 100644 --- a/examples/cpp-dummy/micro_cpp_dummy.cpp +++ b/examples/cpp-dummy/micro_cpp_dummy.cpp @@ -59,6 +59,12 @@ py::list MicroSimulation::get_state() const return state_python; } +// This function needs to return the global id received during construction +int MicroSimulation::get_global_id() const +{ + return _sim_id; +} + PYBIND11_MODULE(micro_dummy, m) { // optional docstring m.doc() = "pybind11 micro dummy plugin"; @@ -68,6 +74,7 @@ PYBIND11_MODULE(micro_dummy, m) { .def("solve", &MicroSimulation::solve) .def("get_state", &MicroSimulation::get_state) .def("set_state", &MicroSimulation::set_state) + .def("get_global_id", &MicroSimulation::get_global_id) // Pickling support does not work currently, as there is no way to pass the simulation ID to the new instance ms. .def(py::pickle( // https://pybind11.readthedocs.io/en/latest/advanced/classes.html#pickling-support [](const MicroSimulation &ms) { // __getstate__ diff --git a/examples/cpp-dummy/micro_cpp_dummy.hpp b/examples/cpp-dummy/micro_cpp_dummy.hpp index fb230ea1..29c306e0 100644 --- a/examples/cpp-dummy/micro_cpp_dummy.hpp +++ b/examples/cpp-dummy/micro_cpp_dummy.hpp @@ -20,6 +20,7 @@ class MicroSimulation void set_state(py::list state); py::list get_state() const; + int get_global_id() const; private: int _sim_id; diff --git a/examples/python-dummy/micro_dummy.py b/examples/python-dummy/micro_dummy.py index f2369f4e..e1e24233 100644 --- a/examples/python-dummy/micro_dummy.py +++ b/examples/python-dummy/micro_dummy.py @@ -1,10 +1,13 @@ """ Micro simulation -In this script we solve a dummy micro problem to just show the working of the macro-micro coupling +In this script we solve a dummy micro problem to just show the working of the macro-micro coupling. +This example shows how to inherit from MicroSimulationInterface provided by the Micro Manager. """ +from micro_manager import MicroSimulationInterface -class MicroSimulation: + +class MicroSimulation(MicroSimulationInterface): def __init__(self, sim_id): """ Constructor of MicroSimulation class. @@ -15,6 +18,9 @@ def __init__(self, sim_id): self._micro_vector_data = None self._state = None + def initialize(self, initial_data=None): + pass + def solve(self, macro_data, dt): assert dt != 0 self._micro_vector_data = [] @@ -32,3 +38,12 @@ def set_state(self, state): def get_state(self): return self._state + + def get_global_id(self): + return self._sim_id + + def set_global_id(self, global_id): + self._sim_id = global_id + + def output(self): + pass diff --git a/micro_manager/__init__.py b/micro_manager/__init__.py index 97a05f31..60a0d9ca 100644 --- a/micro_manager/__init__.py +++ b/micro_manager/__init__.py @@ -1,26 +1,94 @@ import argparse import os -from .config import Config -from .micro_manager import MicroManagerCoupling -try: - from .snapshot.snapshot import MicroManagerSnapshot +def _check_dependencies(): + import importlib.metadata - is_snapshot_possible = True -except ImportError: - is_snapshot_possible = False + from packaging.requirements import Requirement + from packaging.version import Version + _import_name_map = {"pyprecice": "precice"} + required = {} + _pkg_requires = importlib.metadata.requires("micro-manager-precice") or [] + for _dep in _pkg_requires: + if "; extra ==" in _dep: + continue + _req = Requirement(_dep) + _import_name = _import_name_map.get(_req.name, _req.name) + _min_version = None + for _spec in _req.specifier: + if _spec.operator == ">=": + _min_version = _spec.version + required[_req.name] = (_import_name, _min_version) -def main(): + missing = [] + version_errors = [] + + for package, (import_name, min_version) in required.items(): + try: + __import__(import_name) + if min_version: + installed_version = importlib.metadata.version(package) + if Version(installed_version) < Version(min_version): + version_errors.append( + "{} (installed: {}, required: >={})".format( + package, installed_version, min_version + ) + ) + except ImportError: + missing.append(package) + + errors = [] + if missing: + errors.append( + "Missing packages: {}. Install via: pip install {}".format( + ", ".join(missing), " ".join(missing) + ) + ) + if version_errors: + errors.append( + "Version requirements not met: {}".format(", ".join(version_errors)) + ) + + if errors: + raise ImportError("\n".join(errors)) + + print("All dependencies are correctly installed.") + + +import sys + +# Delay heavy imports if only running dependency check +if "--test-dependencies" not in sys.argv: + from .config import Config + from .micro_simulation import MicroSimulationInterface + from .micro_manager import MicroManagerCoupling + + try: + from .snapshot.snapshot import MicroManagerSnapshot + + is_snapshot_possible = True + except ImportError: + is_snapshot_possible = False + +def main(): parser = argparse.ArgumentParser(description=".") parser.add_argument( - "config_file", type=str, help="Path to the JSON config file of the manager." + "config_file", + type=str, + nargs="?", + help="Path to the JSON config file of the manager.", ) parser.add_argument( "--snapshot", action="store_true", help="compute offline snapshot database" ) + parser.add_argument( + "--test-dependencies", + action="store_true", + help="Check if all required dependencies are correctly installed.", + ) parser.add_argument( "log_file", type=str, @@ -30,6 +98,14 @@ def main(): ) args = parser.parse_args() + + if args.test_dependencies: + _check_dependencies() + return + + if not args.config_file: + parser.error("config_file is required unless --test-dependencies is used.") + config_file_path = args.config_file if not os.path.isabs(config_file_path): config_file_path = os.getcwd() + "/" + config_file_path @@ -45,5 +121,4 @@ def main(): ) manager.initialize() - manager.solve() diff --git a/micro_manager/adaptivity/adaptivity.py b/micro_manager/adaptivity/adaptivity.py index 293edb7f..edfb8897 100644 --- a/micro_manager/adaptivity/adaptivity.py +++ b/micro_manager/adaptivity/adaptivity.py @@ -1,17 +1,26 @@ """ Functionality for adaptive initialization and control of micro simulations """ + from math import exp from typing import Callable -import importlib from micro_manager.tools.logging_wrapper import Logger +from micro_manager.config import Config +from micro_manager.micro_simulation import MicroSimulationClass +from micro_manager.model_manager import ModelManager import numpy as np class AdaptivityCalculator: def __init__( - self, configurator, nsims, micro_problem_cls, base_logger, rank + self, + configurator: Config, + nsims: int, + micro_problem_cls: MicroSimulationClass, + model_manager: ModelManager, + base_logger: Logger, + rank: int, ) -> None: """ Class constructor. @@ -24,6 +33,8 @@ def __init__( Number of micro simulations. micro_problem_cls : callable Class of micro problem. + model_manager : object + Handles instantiation of micro simulation. base_logger : object of class Logger Logger object to log messages. rank : int @@ -37,6 +48,7 @@ def __init__( self._adaptivity_output_type = configurator.get_adaptivity_output_type() self._micro_problem_cls = micro_problem_cls + self._model_manager = model_manager self._coarse_tol = 0.0 self._ref_tol = 0.0 @@ -112,7 +124,7 @@ def _update_similarity_dists(self, dt: float, data: dict) -> None: self._similarity_dists *= exp(-self._hist_param * dt) for name in data.keys(): - data_vals = np.array(data[name]) + data_vals = np.asarray(data[name]) if data_vals.ndim == 1: # If the adaptivity data is a scalar for each simulation, # expand the dimension to make it a 2D array to unify the calculation. @@ -183,7 +195,7 @@ def _check_for_deactivation(self, active_id: int, active_ids: list) -> bool: for active_id_2 in active_ids: if active_id != active_id_2: # don't compare active sim to itself # If active sim is similar to another active sim, deactivate it - if self._similarity_dists[active_id, active_id_2] < self._coarse_tol: + if self._similarity_dists[active_id, active_id_2] <= self._coarse_tol: return True return False @@ -263,18 +275,15 @@ def _l1rel(self, data: np.ndarray) -> np.ndarray: similarity_dists : numpy array Updated 2D array having similarity distances between each micro simulation pair """ - pointwise_diff = data[np.newaxis, :] - data[:, np.newaxis] - # divide by data to get relative difference - # divide i,j by max(abs(data[i]),abs(data[j])) to get relative difference - relative = np.nan_to_num( - ( - pointwise_diff - / np.maximum( - np.absolute(data[np.newaxis, :]), np.absolute(data[:, np.newaxis]) - ) - ) + eps = np.finfo(np.float64).eps + data_bc = data[np.newaxis, :] + data_abs = np.absolute(data_bc) + denom = np.maximum(data_abs, np.swapaxes(data_abs, 0, 1)) + return np.linalg.norm( + (data_bc - np.swapaxes(data_bc, 0, 1)) / np.maximum(denom, eps), + ord=1, + axis=-1, ) - return np.linalg.norm(relative, ord=1, axis=-1) def _l2rel(self, data: np.ndarray) -> np.ndarray: """ @@ -291,15 +300,12 @@ def _l2rel(self, data: np.ndarray) -> np.ndarray: similarity_dists : numpy array Updated 2D array having similarity distances between each micro simulation pair """ - pointwise_diff = data[np.newaxis, :] - data[:, np.newaxis] - # divide by data to get relative difference - # divide i,j by max(abs(data[i]),abs(data[j])) to get relative difference - relative = np.nan_to_num( - ( - pointwise_diff - / np.maximum( - np.absolute(data[np.newaxis, :]), np.absolute(data[:, np.newaxis]) - ) - ) + eps = np.finfo(np.float64).eps + data_bc = data[np.newaxis, :] + data_abs = np.absolute(data_bc) + denom = np.maximum(data_abs, np.swapaxes(data_abs, 0, 1)) + return np.linalg.norm( + (data_bc - np.swapaxes(data_bc, 0, 1)) / np.maximum(denom, eps), + ord=2, + axis=-1, ) - return np.linalg.norm(relative, ord=2, axis=-1) diff --git a/micro_manager/adaptivity/adaptivity_selection.py b/micro_manager/adaptivity/adaptivity_selection.py new file mode 100644 index 00000000..a772483c --- /dev/null +++ b/micro_manager/adaptivity/adaptivity_selection.py @@ -0,0 +1,44 @@ +from .global_adaptivity import GlobalAdaptivityCalculator +from .local_adaptivity import LocalAdaptivityCalculator +from .adaptivity import AdaptivityCalculator + + +def create_adaptivity_calculator( + config, + local_number_of_sims, + global_number_of_sims, + global_ids_of_local_sims, + participant, + logger, + rank, + comm, + micro_problem_cls, + model_manager, +) -> AdaptivityCalculator: + adaptivity_type = config.get_adaptivity_type() + + if adaptivity_type == "local": + return LocalAdaptivityCalculator( + config, + local_number_of_sims, + logger, + rank, + comm, + micro_problem_cls, + model_manager, + ) + + if adaptivity_type == "global": + return GlobalAdaptivityCalculator( + config, + global_number_of_sims, + global_ids_of_local_sims, + participant, + logger, + rank, + comm, + micro_problem_cls, + model_manager, + ) + + raise ValueError("Unknown adaptivity type") diff --git a/micro_manager/adaptivity/global_adaptivity.py b/micro_manager/adaptivity/global_adaptivity.py index 81b29a8e..0657443d 100644 --- a/micro_manager/adaptivity/global_adaptivity.py +++ b/micro_manager/adaptivity/global_adaptivity.py @@ -5,27 +5,32 @@ Note: All ID variables used in the methods of this class are global IDs, unless they have *local* in their name. """ -import hashlib from copy import deepcopy -import sys from typing import Dict import numpy as np from mpi4py import MPI from .adaptivity import AdaptivityCalculator +from micro_manager.config import Config +from micro_manager.tools.logging_wrapper import Logger +from micro_manager.micro_simulation import MicroSimulationClass +from micro_manager.model_manager import ModelManager + +from micro_manager.tools.p2p import p2p_comm, get_ranks_of_sims class GlobalAdaptivityCalculator(AdaptivityCalculator): def __init__( self, - configurator, + configurator: Config, global_number_of_sims: int, global_ids: list, participant, - base_logger, + base_logger: Logger, rank: int, - comm, - micro_problem_cls, + comm: MPI.Comm, + micro_problem_cls: MicroSimulationClass, + model_manager: ModelManager, ) -> None: """ Class constructor. @@ -48,15 +53,22 @@ def __init__( Communicator for MPI. micro_problem_cls : callable Class of micro problem. + model_manager : object of class ModelManager + Handles instantiation of the micro simulation. """ super().__init__( - configurator, global_number_of_sims, micro_problem_cls, base_logger, rank + configurator, + global_number_of_sims, + micro_problem_cls, + model_manager, + base_logger, + rank, ) self._global_number_of_sims = global_number_of_sims self._global_ids = global_ids self._comm = comm - rank_of_sim = self._get_ranks_of_sims() + rank_of_sim = get_ranks_of_sims(global_ids, rank, comm, global_number_of_sims) self._is_sim_on_this_rank = [False] * global_number_of_sims # DECLARATION for gid in range(global_number_of_sims): @@ -253,6 +265,12 @@ def get_full_field_micro_output(self, micro_output: list) -> list: return micro_sims_output + def set_is_on_rank(self, gid, val): + """ + Marks whether the simulation gid is as on local rank or not. + """ + self._is_sim_on_this_rank[gid] = val + def log_metrics(self, n: int) -> None: """ Log the following metrics: @@ -289,7 +307,9 @@ def log_metrics(self, n: int) -> None: self._adaptivity_output_type == "all" or self._adaptivity_output_type == "local" ): - ranks_of_sims = self._get_ranks_of_sims() + ranks_of_sims = get_ranks_of_sims( + self._global_ids, self._rank, self._comm, self._global_number_of_sims + ) assoc_ranks = [] # Ranks to which inactive sims on this rank are associated for gid in self._global_ids: @@ -338,15 +358,9 @@ def _update_active_sims(self) -> None: Update set of active micro simulations. Pairs of active simulations (A, B) are compared and if found to be similar, B is deactivated. """ - if self._max_similarity_dist == 0.0: - self._base_logger.log_warning( - "All similarity distances are zero, which means all the data for adaptivity is the same. Coarsening tolerance will be manually set to minimum float number." - ) - self._coarse_tol = sys.float_info.min - else: - self._coarse_tol = ( - self._coarse_const * self._refine_const * self._max_similarity_dist - ) + self._coarse_tol = ( + self._coarse_const * self._refine_const * self._max_similarity_dist + ) active_gids_this_rank = self.get_active_sim_global_ids() # Gather global ids of active sims from all ranks @@ -406,7 +420,15 @@ def _communicate_micro_output( assoc_active_gids = list(active_to_inactive_map.keys()) - recv_reqs = self._p2p_comm(assoc_active_gids, micro_output) + recv_reqs = p2p_comm( + self._global_ids, + self._rank, + self._comm, + self._global_number_of_sims, + self._is_sim_on_this_rank, + assoc_active_gids, + micro_output, + ) # Add received output of active sims to inactive sims on this rank for count, req in enumerate(recv_reqs): @@ -459,7 +481,9 @@ def _update_inactive_sims(self, micro_sims: list) -> None: # Only handle activation of simulations on this rank for gid in to_be_activated_gids: to_be_activated_lid = self._global_ids.index(gid) - micro_sims[to_be_activated_lid] = self._micro_problem_cls(gid) + micro_sims[to_be_activated_lid] = self._model_manager.get_instance( + gid, self._micro_problem_cls + ) assoc_active_gid = self._sim_is_associated_to[gid] if self._is_sim_on_this_rank[ @@ -481,13 +505,19 @@ def _update_inactive_sims(self, micro_sims: list) -> None: sim_states_and_global_ids = [] for lid, sim in enumerate(micro_sims): - if sim == 0: + if sim == 0 or sim is None: sim_states_and_global_ids.append((None, self._global_ids[lid])) else: sim_states_and_global_ids.append((sim.get_state(), sim.get_global_id())) - recv_reqs = self._p2p_comm( - list(to_be_activated_map.keys()), sim_states_and_global_ids + recv_reqs = p2p_comm( + self._global_ids, + self._rank, + self._comm, + self._global_number_of_sims, + self._is_sim_on_this_rank, + list(to_be_activated_map.keys()), + sim_states_and_global_ids, ) # Use received micro sims to activate the required simulations @@ -496,134 +526,23 @@ def _update_inactive_sims(self, micro_sims: list) -> None: local_ids = to_be_activated_map[gid] for lid in local_ids: # Create the micro simulation object and set its state - micro_sims[lid] = self._micro_problem_cls(self._global_ids[lid]) + micro_sims[lid] = self._model_manager.get_instance( + self._global_ids[lid], self._micro_problem_cls + ) micro_sims[lid].set_state(state) # Delete the micro simulation object if it is inactive for gid in self._global_ids: if not self._is_sim_active[gid]: lid = self._global_ids.index(gid) - micro_sims[lid] = 0 + # Release resources now, especially for remote simulation instance. + # If left to garbage collector this might lead to a race condition. + # Releasing with call to sim.destroy(), afterwards reference in sim + # sim list can be removed. + if type(micro_sims[lid]).__name__ == "MicroSimulationWrapper": + micro_sims[lid].destroy() + micro_sims[lid] = None self._precice_participant.stop_last_profiling_section() self._sim_is_associated_to = np.copy(_sim_is_associated_to_updated) - - def _create_tag(self, sim_id: int, src_rank: int, dest_rank: int) -> int: - """ - For a given simulations ID, source rank, and destination rank, a unique tag is created. - - Parameters - ---------- - sim_id : int - Global ID of a simulation. - src_rank : int - Rank on which the simulation lives - dest_rank : int - Rank to which data of a simulation is to be sent to. - - Returns - ------- - tag : int - Unique tag. - """ - send_hashtag = hashlib.sha256() - send_hashtag.update( - (str(src_rank) + str(sim_id) + str(dest_rank)).encode("utf-8") - ) - tag = int(send_hashtag.hexdigest()[:6], base=16) - return tag - - def _p2p_comm(self, assoc_active_ids: list, data: list) -> list: - """ - Handle process to process communication for a given set of associated active IDs and data. - - Parameters - ---------- - assoc_active_ids : list - Global IDs of active simulations which are not on this rank and are associated to - the inactive simulations on this rank. - data : list - Complete data from which parts are to be sent and received. - - Returns - ------- - recv_reqs : list - List of MPI requests of receive operations. - """ - rank_of_sim = self._get_ranks_of_sims() - - send_map_local: Dict[ - int, int - ] = dict() # keys are global IDs, values are rank to send to - send_map: Dict[ - int, list - ] = ( - dict() - ) # keys are global IDs of sims to send, values are ranks to send the sims to - recv_map: Dict[ - int, int - ] = dict() # keys are global IDs to receive, values are ranks to receive from - - for i in assoc_active_ids: - # Add simulation and its rank to receive map - recv_map[i] = rank_of_sim[i] - # Add simulation and this rank to local sending map - send_map_local[i] = self._rank - - # Gather information about which sims to send where, from the sending perspective - send_map_list = self._comm.allgather(send_map_local) - - for d in send_map_list: - for i, rank in d.items(): - if self._is_sim_on_this_rank[i]: - if i in send_map: - send_map[i].append(rank) - else: - send_map[i] = [rank] - - # Asynchronous send operations - send_reqs = [] - for gid, send_ranks in send_map.items(): - lid = self._global_ids.index(gid) - for send_rank in send_ranks: - tag = self._create_tag(gid, self._rank, send_rank) - req = self._comm.isend(data[lid], dest=send_rank, tag=tag) - send_reqs.append(req) - - # Asynchronous receive operations - recv_reqs = [] - for gid, recv_rank in recv_map.items(): - tag = self._create_tag(gid, recv_rank, self._rank) - bufsize = ( - 1 << 30 - ) # allocate and use a temporary 1 MiB buffer size https://github.com/mpi4py/mpi4py/issues/389 - req = self._comm.irecv(bufsize, source=recv_rank, tag=tag) - recv_reqs.append(req) - - # Wait for all non-blocking communication to complete - MPI.Request.Waitall(send_reqs) - - return recv_reqs - - def _get_ranks_of_sims(self) -> np.ndarray: - """ - Get the ranks of all simulations. - - Returns - ------- - ranks_of_sims : np.ndarray - Array of ranks on which simulations exist. - """ - gids_to_rank = dict() - for gid in self._global_ids: - gids_to_rank[gid] = self._rank - - ranks_maps_as_list = self._comm.allgather(gids_to_rank) - - ranks_of_sims = np.zeros(self._global_number_of_sims, dtype=np.intc) - for ranks_map in ranks_maps_as_list: - for gid, rank in ranks_map.items(): - ranks_of_sims[gid] = rank - - return ranks_of_sims diff --git a/micro_manager/adaptivity/global_adaptivity_lb.py b/micro_manager/adaptivity/global_adaptivity_lb.py deleted file mode 100644 index e6ded12f..00000000 --- a/micro_manager/adaptivity/global_adaptivity_lb.py +++ /dev/null @@ -1,423 +0,0 @@ -""" -Class GlobalAdaptivityLBCalculator provides methods to adaptively control of micro simulations -in a global way. If the Micro Manager is run in parallel, an all-to-all comparison of simulations -on each rank is done, along with dynamic load balancing. - -Note: All ID variables used in the methods of this class are global IDs, unless they have *local* in their name. -""" -import numpy as np -from mpi4py import MPI -import math - -from .global_adaptivity import GlobalAdaptivityCalculator - - -class GlobalAdaptivityLBCalculator(GlobalAdaptivityCalculator): - def __init__( - self, - configurator, - global_number_of_sims: int, - global_ids: list, - participant, - base_logger, - rank: int, - comm, - micro_problem_cls: callable, - ) -> None: - """ - Class constructor. - - Parameters - ---------- - configurator : object of class Config - Object which has getter functions to get parameters defined in the configuration file. - global_number_of_sims : int - Total number of simulations in the macro-micro coupled problem. - global_ids : list - List of global IDs of simulations living on this rank. - participant : object of class Participant - Object of the class Participant using which the preCICE API is called. - base_logger : object of class Logger - Logger to log to terminal. - rank : int - MPI rank. - comm : MPI.Comm - Communicator for MPI. - micro_problem_cls : callable - Class of micro problem. - """ - super().__init__( - configurator, - global_number_of_sims, - global_ids, - participant, - base_logger, - rank, - comm, - micro_problem_cls, - ) - - self._base_logger = base_logger - - self._threshold = configurator.get_load_balancing_threshold() - - self._balance_inactive_sims = configurator.balance_inactive_sims() - - self._nothing_to_balance = False - - self._precice_participant = participant - - def redistribute_sims(self, micro_sims: list) -> None: - """ - Redistribute simulations among ranks to balance compute load. - - Parameters - ---------- - micro_sims : list - List of objects of class MicroProblem, which are the micro simulations - """ - self._nothing_to_balance = False - - self._precice_participant.start_profiling_section( - "global_adaptivity_lb.redistributing_sims" - ) - - self._redistribute_active_sims(micro_sims) - - if (not self._nothing_to_balance) and self._balance_inactive_sims: - self._redistribute_inactive_sims(micro_sims) - - self._precice_participant.stop_last_profiling_section() - - def _redistribute_active_sims(self, micro_sims: list) -> None: - """ - Redistribute active simulations as per the configured load balancing settings. - - Parameters - ---------- - micro_sims : list - List of objects of class MicroProblem, which are the micro simulations - """ - avg_active_sims = np.count_nonzero(self._is_sim_active) / self._comm.size - - active_sims_local_ids = self.get_active_sim_local_ids() - - n_active_sims_local = active_sims_local_ids.size - - send_sims = 0 # Sims that this rank wants to send - recv_sims = 0 # Sims that this rank wants to receive - - f_avg_active_sims = math.floor(avg_active_sims - self._threshold) - c_avg_active_sims = math.ceil(avg_active_sims + self._threshold) - - if f_avg_active_sims == c_avg_active_sims: - if n_active_sims_local < avg_active_sims: - recv_sims = int(avg_active_sims) - n_active_sims_local - elif n_active_sims_local > avg_active_sims: - send_sims = n_active_sims_local - int(avg_active_sims) - else: - if n_active_sims_local < f_avg_active_sims: - recv_sims = f_avg_active_sims - n_active_sims_local - elif n_active_sims_local == f_avg_active_sims: - recv_sims += 1 - elif n_active_sims_local > c_avg_active_sims: - send_sims = n_active_sims_local - c_avg_active_sims - elif n_active_sims_local == c_avg_active_sims: - send_sims += 1 - - # Number of active sims that each rank wants to send and receive - global_send_sims = self._comm.allgather(send_sims) - global_recv_sims = self._comm.allgather(recv_sims) - - n_global_send_sims = sum(global_send_sims) - n_global_recv_sims = sum(global_recv_sims) - - if n_global_send_sims == 0 and n_global_recv_sims == 0: - self._nothing_to_balance = True - self._base_logger.log_warning_rank_zero( - "It appears that the micro simulations are already fairly balanced. No load balancing will be done. Try changing the threshold value to induce load balancing." - ) - return - - if n_global_send_sims < n_global_recv_sims: - excess_recv_sims = n_global_recv_sims - n_global_send_sims - while excess_recv_sims > 0: - for i, e in enumerate(global_recv_sims): - if e > 0: - # Remove the excess receive request from the rank - global_recv_sims[i] -= 1 - - excess_recv_sims -= 1 - - if excess_recv_sims == 0: - break - elif n_global_send_sims > n_global_recv_sims: - excess_send_sims = n_global_send_sims - n_global_recv_sims - while excess_send_sims > 0: - for i, e in enumerate(global_send_sims): - if e > 0: - # Remove the excess send request - global_send_sims[i] -= 1 - - excess_send_sims -= 1 - - if excess_send_sims == 0: - break - - send_map, recv_map = self._get_communication_maps( - global_send_sims, global_recv_sims - ) - - self._move_active_sims(micro_sims, send_map, recv_map) - - def _redistribute_inactive_sims(self, micro_sims: list) -> None: - """ - Redistribute inactive simulations based on where the associated active simulations are. - - Parameters - ---------- - micro_sims : list - List of objects of class MicroProblem, which are the micro simulations - """ - # Dict of - # keys: global IDs of sim states to send from this rank - # values: ranks to send the sims to - send_map: dict[int, int] = dict() - - # Dict of - # keys: global IDs of sim states to receive on this rank - # values: are ranks to receive from - recv_map: dict[int, int] = dict() - - ranks_of_sims = self._get_ranks_of_sims() - - global_ids_of_inactive_sims = np.where(self._is_sim_active == False)[0] - - for gid in global_ids_of_inactive_sims: - assoc_active_gid = self._sim_is_associated_to[gid] - rank_of_inactive_sim = ranks_of_sims[gid] - rank_of_assoc_active_sim = ranks_of_sims[assoc_active_gid] - if rank_of_inactive_sim != rank_of_assoc_active_sim: - if rank_of_inactive_sim == self._rank: - send_map[gid] = rank_of_assoc_active_sim - if rank_of_assoc_active_sim == self._rank: - recv_map[gid] = rank_of_inactive_sim - - self._move_inactive_sims(micro_sims, send_map, recv_map) - - def _get_communication_maps( - self, global_send_sims: list, global_recv_sims: list - ) -> tuple: - """ - Create dictionaries which map global IDs of simulations to ranks for sending and receiving. - - Parameters - ---------- - global_send_sims : list - Number of simulations that each rank sends. - global_recv_sims : list - Number of simulations that each rank receives. - - Returns - ------- - tuple of dicts - send_map : dict - keys are global IDs of sim states to send, values are ranks to send the sims to - recv_map : dict - keys are global IDs of sim states to receive, values are ranks to receive from - """ - active_sims_global_ids = list(self.get_active_sim_global_ids()) - - rank_wise_global_ids_of_active_sims = self._comm.allgather( - active_sims_global_ids - ) - - # Keys are ranks sending sims, values are lists of tuples: (list of global IDs to send, the rank to send them to) - global_send_map: dict[int, list] = dict() - - # Keys are ranks receiving sims, values are lists of tuples: (list of global IDs to receive, the rank to receive them from) - global_recv_map: dict[int, list] = dict() - - for rank in [i for i, e in enumerate(global_send_sims) if e != 0]: - global_send_map[rank] = [] - - for rank in [i for i, e in enumerate(global_recv_sims) if e != 0]: - global_recv_map[rank] = [] - - send_ranks = list(global_send_map.keys()) - recv_ranks = list(global_recv_map.keys()) - - count = 0 - recv_rank = recv_ranks[count] - - for send_rank in send_ranks: - sims = global_send_sims[send_rank] - while sims > 0: - if global_recv_sims[recv_rank] <= sims: - # Get the global IDs to move - global_ids_of_sims_to_move = rank_wise_global_ids_of_active_sims[ - send_rank - ][0 : int(global_recv_sims[recv_rank])] - - global_send_map[send_rank].append( - (global_ids_of_sims_to_move, recv_rank) - ) - - global_recv_map[recv_rank].append( - (global_ids_of_sims_to_move, send_rank) - ) - - sims -= global_recv_sims[recv_rank] - - # Remove the global IDs which are already mapped for moving - del rank_wise_global_ids_of_active_sims[send_rank][ - 0 : int(global_recv_sims[recv_rank]) - ] - - if count < len(recv_ranks) - 1: - count += 1 - recv_rank = recv_ranks[count] - - elif global_recv_sims[recv_rank] > sims: - # Get the global IDs to move - global_ids_of_sims_to_move = rank_wise_global_ids_of_active_sims[ - send_rank - ][0 : int(sims)] - - global_send_map[send_rank].append( - (global_ids_of_sims_to_move, recv_rank) - ) - - global_recv_map[recv_rank].append( - (global_ids_of_sims_to_move, send_rank) - ) - - global_recv_sims[recv_rank] -= sims - - # Remove the global IDs which are already mapped for moving - del rank_wise_global_ids_of_active_sims[send_rank][0 : int(sims)] - - sims = 0 - - # keys are global IDs of sim states to send, values are ranks to send the sims to - send_map: dict[int, int] = dict() - - # keys are global IDs of sim states to receive, values are ranks to receive from - recv_map: dict[int, int] = dict() - - if self._rank in global_send_map: - for send_info in global_send_map[self._rank]: - send_rank = send_info[1] - for gid in send_info[0]: - send_map[gid] = send_rank - - if self._rank in global_recv_map: - for recv_info in global_recv_map[self._rank]: - recv_rank = recv_info[1] - for gid in recv_info[0]: - recv_map[gid] = recv_rank - - return send_map, recv_map - - def _move_active_sims( - self, micro_sims: list, send_map: dict, recv_map: dict - ) -> None: - """ - Move active micro simulations between ranks. - - Parameters - ---------- - micro_sims : list - List of objects of class MicroProblem, which are the micro simulations - send_map : dict - keys are global IDs of sim states to send, values are ranks to send the sims to - recv_map : dict - keys are global IDs of sim states to receive, values are ranks to receive from - """ - # Asynchronous send operations - send_reqs = [] - for gid, send_rank in send_map.items(): - tag = self._create_tag(gid, self._rank, send_rank) - lid = self._global_ids.index(gid) - req = self._comm.isend( - (micro_sims[lid].get_state(), gid), dest=send_rank, tag=tag - ) - send_reqs.append(req) - - # Asynchronous receive operations - recv_reqs = [] - for gid, recv_rank in recv_map.items(): - tag = self._create_tag(gid, recv_rank, self._rank) - bufsize = ( - 1 << 30 - ) # allocate and use a temporary 1 MiB buffer size https://github.com/mpi4py/mpi4py/issues/389 - req = self._comm.irecv(bufsize, source=recv_rank, tag=tag) - recv_reqs.append(req) - - # Wait for all non-blocking communication to complete - MPI.Request.Waitall(send_reqs) - - # Delete the active simulations which no longer exist on this rank - for gid in send_map.keys(): - lid = self._global_ids.index(gid) - del micro_sims[lid] - self._global_ids.remove(gid) - self._is_sim_on_this_rank[gid] = False - - # Create simulations and set them to the received states - for req in recv_reqs: - output, gid = req.wait() - micro_sims.append(self._micro_problem_cls(gid)) - micro_sims[-1].set_state(output) - self._global_ids.append(gid) - self._is_sim_on_this_rank[gid] = True - - def _move_inactive_sims( - self, micro_sims: list, send_map: dict, recv_map: dict - ) -> None: - """ - Move inactive micro simulation states between ranks. - - Parameters - ---------- - micro_sims : list - List of objects of class MicroProblem, which are the micro simulations - send_map : dict - keys are global IDs of sim states to send, values are ranks to send the sims to - recv_map : dict - keys are global IDs of sim states to receive, values are ranks to receive from - """ - # Asynchronous send operations - send_reqs = [] - for gid, send_rank in send_map.items(): - tag = self._create_tag(gid, self._rank, send_rank) - lid = self._global_ids.index(gid) - req = self._comm.isend((gid), dest=send_rank, tag=tag) - send_reqs.append(req) - - # Asynchronous receive operations - recv_reqs = [] - for gid, recv_rank in recv_map.items(): - tag = self._create_tag(gid, recv_rank, self._rank) - bufsize = ( - 1 << 30 - ) # allocate and use a temporary 1 MiB buffer size https://github.com/mpi4py/mpi4py/issues/389 - req = self._comm.irecv(bufsize, source=recv_rank, tag=tag) - recv_reqs.append(req) - - # Wait for all non-blocking communication to complete - MPI.Request.Waitall(send_reqs) - - # Delete the inactive simulations which no longer exist on this rank - for gid in send_map.keys(): - lid = self._global_ids.index(gid) - del micro_sims[lid] - self._global_ids.remove(gid) - self._is_sim_on_this_rank[gid] = False - - # Add inactive simulations in the data structure - for req in recv_reqs: - gid = req.wait() - micro_sims.append(0) - self._global_ids.append(gid) - self._is_sim_on_this_rank[gid] = True diff --git a/micro_manager/adaptivity/local_adaptivity.py b/micro_manager/adaptivity/local_adaptivity.py index 6ea9a717..bbd9d614 100644 --- a/micro_manager/adaptivity/local_adaptivity.py +++ b/micro_manager/adaptivity/local_adaptivity.py @@ -3,23 +3,27 @@ in a local way. If the Micro Manager is run in parallel, simulations on one rank are compared to each other. A global comparison is not done. """ -import sys import numpy as np from copy import deepcopy from mpi4py import MPI from .adaptivity import AdaptivityCalculator +from micro_manager.config import Config +from micro_manager.micro_simulation import MicroSimulationClass +from micro_manager.tools.logging_wrapper import Logger +from micro_manager.model_manager import ModelManager class LocalAdaptivityCalculator(AdaptivityCalculator): def __init__( self, - configurator, - num_sims, - base_logger, - rank, - comm, - micro_problem_cls, + configurator: Config, + num_sims: int, + base_logger: Logger, + rank: int, + comm: MPI.Comm, + micro_problem_cls: MicroSimulationClass, + model_manager: ModelManager, ) -> None: """ Class constructor. @@ -38,8 +42,12 @@ def __init__( Communicator for MPI. micro_problem_cls : callable Class of micro problem. + model_manager : object of class ModelManager + Handles instantiation of micro simulation. """ - super().__init__(configurator, num_sims, micro_problem_cls, base_logger, rank) + super().__init__( + configurator, num_sims, micro_problem_cls, model_manager, base_logger, rank + ) self._comm = comm # similarity_dists: 2D array having similarity distances between each micro simulation pair @@ -48,8 +56,8 @@ def __init__( def compute_adaptivity( self, - dt, - micro_sims, + dt: float, + micro_sims: list, data_for_adaptivity: dict, ) -> None: """ @@ -237,15 +245,9 @@ def _update_active_sims(self) -> None: Update set of active micro simulations. Active micro simulations are compared to each other and if found similar, one of them is deactivated. """ - if self._max_similarity_dist == 0.0: - self._base_logger.log_warning( - "All similarity distances are zero, which means all the data for adaptivity is the same. Coarsening tolerance will be manually set to minimum float number." - ) - self._coarse_tol = sys.float_info.min - else: - self._coarse_tol = ( - self._coarse_const * self._refine_const * self._max_similarity_dist - ) + self._coarse_tol = ( + self._coarse_const * self._refine_const * self._max_similarity_dist + ) active_gids = self.get_active_sim_local_ids().tolist() @@ -292,7 +294,7 @@ def _update_inactive_sims(self, micro_sims: list) -> None: # Update the set of inactive micro sims for i in to_be_activated_ids: associated_active_id = self._sim_is_associated_to[i] - micro_sims[i] = self._micro_problem_cls(i) + micro_sims[i] = self._model_manager.get_instance(i, self._micro_problem_cls) micro_sims[i].set_state(micro_sims[associated_active_id].get_state()) self._sim_is_associated_to[ i @@ -301,4 +303,10 @@ def _update_inactive_sims(self, micro_sims: list) -> None: # Delete the inactive micro simulations which have not been activated for i in range(self._is_sim_active.size): if not self._is_sim_active[i]: - micro_sims[i] = 0 + # Release resources now, especially for remote simulation instance. + # If left to garbage collector this might lead to a race condition. + # Releasing with call to sim.destroy(), afterwards reference in sim + # sim list can be removed. + if type(micro_sims[i]).__name__ == "MicroSimulationWrapper": + micro_sims[i].destroy() + micro_sims[i] = None diff --git a/micro_manager/adaptivity/model_adaptivity.py b/micro_manager/adaptivity/model_adaptivity.py index 7d4e049f..a0c4b521 100644 --- a/micro_manager/adaptivity/model_adaptivity.py +++ b/micro_manager/adaptivity/model_adaptivity.py @@ -3,45 +3,77 @@ """ from typing import Union, Optional -from ..config import Config -from ..micro_simulation import create_simulation_class +from micro_manager.config import Config +from micro_manager.micro_simulation import ( + create_simulation_class, + load_backend_class, + MicroSimulationClass, +) from micro_manager.tools.logging_wrapper import Logger from micro_manager.tools.misc import clamp_in_range +from micro_manager.model_manager import ModelManager +from micro_manager.tasking.connection import Connection +from mpi4py import MPI import numpy as np import importlib class ModelAdaptivity: - def __init__(self, configurator: Config, rank: int, log_file: str) -> None: + def __init__( + self, + model_manager: ModelManager, + configurator: Config, + comm: MPI.Comm, + rank: int, + log_file: str, + conn: Connection, + num_ranks: int, + ) -> None: """ Class constructor. Parameters ---------- + model_manager: ModelManager + ModelManager instance configurator : object of class Config Object which has getter functions to get parameters defined in the configuration file. + comm: MPI.Comm + MPI communicator rank : int Rank of the MPI communicator. log_file : str Path to the log file to write to. + conn: Connection + Connection to workers + num_ranks : int + Number of workers """ self._logger = Logger(__name__, log_file, rank) + self._comm = comm + self._model_manager = model_manager self._model_files = configurator.get_model_adaptivity_file_names() self._switching_func_name = ( configurator.get_model_adaptivity_switching_function() ) + stateless_flags = configurator.get_model_adaptivity_micro_stateless() self._model_classes = [] - CLASS_NAME = "MicroSimulation" + pos = 0 for model_file in self._model_files: try: - model = getattr( - importlib.import_module(model_file, CLASS_NAME), - CLASS_NAME, + model = load_backend_class(model_file) + self._model_classes.append( + create_simulation_class( + self._logger, model, model_file, num_ranks, conn + ) ) - self._model_classes.append(create_simulation_class(model)) + self._model_manager.register( + self._model_classes[pos], stateless_flags[pos] + ) + pos += 1 except Exception as e: self._logger.log_info_rank_zero( f"Failed to load model class with error: {e}" @@ -84,8 +116,8 @@ def switching_interface( Array with gaussian point for respective sim. D is the mesh dimension. t : float Current time in simulation. - inputs : list[dict] - List of input objects. + input : dict + input object. prev_output : [None, dict-like] Contains the output of the previous model evaluation. @@ -115,10 +147,10 @@ def switch_models( locations: np.ndarray, t: float, inputs: list[dict], - prev_output: dict, + prev_output: Optional[list[dict]], sims: list, - active_sim_ids: Optional[list[int]] = None, - ) -> None: + active_sim_ids: Optional[list] = None, + ) -> list[int]: """ Switches models within sims list. If active_sim_ids is None, all sims are considered as active. @@ -136,33 +168,119 @@ def switch_models( List of all simulation objects. active_sim_ids : [list, None] List of all active simulation ids. + + Returns + ------- + switched_lids : list[int] + List of lids of simulations that were switched. """ size = len(sims) active_sims = self._create_active_mask(active_sim_ids, size) - cur_res = self._gather_current_resolutions(sims, active_sims) - tgt_res = self._gather_target_resolutions( - cur_res, locations, t, inputs, prev_output, active_sims + current_res = self._gather_current_resolutions(sims, active_sims) + target_res = self._gather_target_resolutions( + current_res, locations, t, inputs, prev_output, active_sims ) - self._logger.log_info_rank_zero(f"New resolutions for t={t}: {tgt_res}") + self._logger.log_info(f"New resolutions for t={t}: {target_res}") + + for idx in range(size): + if current_res[idx] == target_res[idx]: + continue + + sim = sims[idx] + gid = sim.get_global_id() + target_class = self.get_resolution_sim_class(target_res[idx]) + + # we store state for each resolution separately + # keys are sim names of respective resolution + key = f"{sim.name}-state" + key_new = f"{target_class.name}-state" + + # check if a state of the target resolution exists + # then update state buffer with current state + new_state_exists = key_new in sim.attachments + sim.attachments[key] = sim.get_state() + + # construct new sim and delay initialization if possible + sim_new = self._model_manager.get_instance( + gid, target_class, late_init=new_state_exists + ) + # need to copy over the multi-state buffer to new sim object + sim_new.attachments = sim.attachments + sim_new.attachments[key_new] = sim_new.get_state() + + # if state of target resolution exists + # use it to initialize + if new_state_exists: + sim_new_state = sim.attachments[key_new] + sim_new.set_state(sim_new_state) + + # release resources of previous sim and set to new sim + sims[idx].destroy() + sims[idx] = sim_new + + return np.argwhere((current_res - target_res) != 0).tolist() + + def update_states( + self, + sims: list, + active_sim_ids: Optional[list] = None, + ): + """ + Updates the current state of the current model in the local buffers. + + Parameters + ---------- + sims : list + List of all simulation objects. + active_sim_ids : [list, None] + List of all active simulation ids. + """ + size = len(sims) + active_sims = self._create_active_mask(active_sim_ids, size) + + for idx in range(size): + if not active_sims[idx]: + continue + + sim = sims[idx] + key = f"{sim.name}-state" + sim.attachments[key] = sim.get_state() + + def write_back_states( + self, + sims: list, + active_sim_ids: Optional[list] = None, + ): + """ + Loads the current state of the current model into local buffers. + + Parameters + ---------- + sims : list + List of all simulation objects. + active_sim_ids : [list, None] + List of all active simulation ids. + """ + size = len(sims) + active_sims = self._create_active_mask(active_sim_ids, size) for idx in range(size): - if cur_res[idx] == tgt_res[idx]: + if not active_sims[idx]: continue - sim_state = sims[idx].get_state() - sim_id = sims[idx].get_global_id() - sims[idx] = self.get_resolution_sim_class(tgt_res[idx])(sim_id) - sims[idx].set_state(sim_state) + sim = sims[idx] + key = f"{sim.name}-state" + sim.set_state(sim.attachments[key]) def check_convergence( self, locations: np.ndarray, t: float, - inputs: list[dict], - prev_output: Optional[dict], + inputs: list, + prev_output: Optional[list[dict]], sims: list, - active_sim_ids: Optional[list[int]] = None, + active_sim_ids: Optional[list] = None, ) -> None: """ Similarly to switch_models, checks whether models would be switched in next step. @@ -194,7 +312,9 @@ def check_convergence( next_switch[idx] = self._switching_func( resolutions[idx], locations[idx], t, inputs[idx], prev_out ) - self._converged = np.all(next_switch == 0) + local_num_changes = np.sum(next_switch != 0) + global_num_changes = self._comm.allreduce(local_num_changes, op=MPI.SUM) + self._converged = global_num_changes == 0 def get_num_resolutions(self) -> int: """ @@ -209,7 +329,7 @@ def get_num_resolutions(self) -> int: def get_resolution_sim_class( self, resolution: Union[int, np.ndarray] - ) -> Union[object, np.ndarray]: + ) -> Union[MicroSimulationClass, list[MicroSimulationClass]]: """ Looks up the class associated with the provided resolution. @@ -227,9 +347,7 @@ def get_resolution_sim_class( clamp_in_range(resolution, 0, len(self._model_classes) - 1) ] - def get_sim_class_resolution( - self, sim: Union[object, np.ndarray] - ) -> Union[int, np.ndarray]: + def get_sim_class_resolution(self, sim: MicroSimulationClass) -> int: """ Looks up the resolution associated with the provided simulation object. @@ -244,11 +362,11 @@ def get_sim_class_resolution( target resolution """ return next( - (idx for idx, cls in enumerate(self._model_classes) if cls == type(sim)) + (idx for idx, cls in enumerate(self._model_classes) if cls.name == sim.name) ) def _gather_current_resolutions( - self, sims: list[object], active_sims: np.ndarray + self, sims: list, active_sims: np.ndarray ) -> np.ndarray: """ Gathers current resolutions. Inactive sims have resolution -1. @@ -278,7 +396,7 @@ def _gather_target_resolutions( locations: np.ndarray, t: float, inputs: list[dict], - prev_output: dict, + prev_output: Optional[list[dict]], active_sims: np.ndarray, ) -> np.ndarray: """ @@ -320,7 +438,7 @@ def _gather_target_resolutions( ) return res_tgt - def _create_active_mask(self, active_sim_ids: list[int], size: int) -> np.ndarray: + def _create_active_mask(self, active_sim_ids: list, size: int) -> np.ndarray: """ Converts list of active simulation ids to np boolean mask. diff --git a/micro_manager/config.py b/micro_manager/config.py index a88b816b..7d6b634b 100644 --- a/micro_manager/config.py +++ b/micro_manager/config.py @@ -25,6 +25,7 @@ def __init__(self, config_file_name): self._config_file_name = config_file_name self._logger = None self._micro_file_name = None + self._micro_stateless = False self._precice_config_file_name = None self._macro_mesh_name = None @@ -32,8 +33,12 @@ def __init__(self, config_file_name): self._write_data_names = None self._micro_dt = None + # Domain decomposition information self._macro_domain_bounds = None self._ranks_per_axis = None + self._decomposition_type = "uniform" + self._minimum_access_region_size: list = [] + self._micro_output_n = 1 self._diagnostics_data_names = None @@ -54,10 +59,12 @@ def __init__(self, config_file_name): self._adaptivity_output_type = "" self._adaptivity_output_n = 1 - self._adaptivity_is_load_balancing = False + self._load_balancing = False + self._load_balancing_type = "time" self._load_balancing_n = 1 + self._load_balancing_partitioning = "lpt" self._load_balancing_threshold = 0 - self._balance_inactive_sims = False + self._load_balancing_balance_inactive_sims = False # Snapshot information self._parameter_file_name = None @@ -72,8 +79,16 @@ def __init__(self, config_file_name): # Model Adaptivity information self._m_adap = False self._m_adap_micro_file_names = None + self._m_adap_micro_stateless = None self._m_adap_switching_function = None + # Tasking + self._task_is_slurm = False + self._task_backend = "socket" + self._task_num_workers = 1 + self._task_mpi_impl = "open" + self._task_pinning_hostfile = "./hosts.micro" + def set_logger(self, logger): """ Set the logger for the Config class. @@ -114,6 +129,17 @@ def _read_json(self, config_file_name): .replace(".py", "") ) + try: + self._micro_stateless = self._data["micro_stateless"] + self._logger.log_info_rank_zero( + "Only creating one full instance of MicroSimulation." + ) + except: + self._micro_stateless = False + self._logger.log_info_rank_zero( + "Creating an instance of MicroSimulation for each mesh vertex." + ) + self._logger.log_info_rank_zero( "Micro simulation file name: " + self._data["micro_file_name"] ) @@ -182,6 +208,29 @@ def _read_json(self, config_file_name): self._micro_dt = self._data["simulation_params"]["micro_dt"] + try: + if self._data["tasking"]: + backend = self._data["tasking"]["backend"] + if backend not in ["mpi", "socket"]: + raise Exception("Backend must be either 'mpi' or 'socket'.") + self._task_backend = backend + if "is_slurm" in self._data["tasking"]: + self._task_is_slurm = self._data["tasking"]["is_slurm"] + if "num_workers" in self._data["tasking"]: + self._task_num_workers = self._data["tasking"]["num_workers"] + if self._task_is_slurm and backend == "mpi": + raise Exception("MPI backend not supported on SLURM systems.") + if "mpi_impl" in self._data["tasking"]: + self._task_mpi_impl = self._data["tasking"]["mpi_impl"] + if self._task_mpi_impl not in ["open", "intel"]: + raise Exception("mpi_impl must be either 'open' or 'intel'.") + if "hostfile" in self._data["tasking"]: + self._task_pinning_hostfile = self._data["tasking"]["hostfile"] + except BaseException: + self._logger.log_info_rank_zero( + "No or incorrect tasking information provided. Micro manager will not create workers and instead solve micro simulations locally." + ) + def read_json_micro_manager(self): """ Reads Micro Manager relevant information from JSON configuration file @@ -213,6 +262,30 @@ def read_json_micro_manager(self): self._logger.log_info_rank_zero( "Axis-wise domain decomposition: " + str(self._ranks_per_axis) ) + if self._data["simulation_params"]["decomposition_type"]: + self._decomposition_type = self._data["simulation_params"][ + "decomposition_type" + ] + if self._decomposition_type not in ["uniform", "nonuniform"]: + raise Exception( + "Decomposition type can be either 'uniform' or 'nonuniform'." + ) + if ( + self._data["simulation_params"]["decomposition_type"] + == "nonuniform" + ): + if self._data["simulation_params"]["minimum_access_region_size"]: + self._minimum_access_region_size = self._data[ + "simulation_params" + ]["minimum_access_region_size"] + else: + self._logger.log_info_rank_zero( + "Minimum access region size is not specified. Calculating it as 1 / (2^ranks_per_axis - 1) of the macro domain size in each axis." + ) + + self._logger.log_info_rank_zero( + "Domain decomposition type: " + self._decomposition_type + ) except BaseException: self._logger.log_info_rank_zero( "Domain decomposition is not specified, so the Micro Manager will expect to be run in serial." @@ -389,24 +462,22 @@ def read_json_micro_manager(self): self._write_data_names.append("Active-Steps") try: - self._adaptivity_is_load_balancing = self._data["simulation_params"][ - "load_balancing" - ] - if self._adaptivity_is_load_balancing: + self._load_balancing = self._data["simulation_params"]["load_balancing"] + if self._load_balancing: self._logger.log_info_rank_zero( - "Micro Manager will dynamically balance micro simulations based on the adaptivity computation." + "Micro Manager will dynamically balance micro simulations." ) self._write_data_names.append("rank_of_sim") - if not self._adaptivity_type == "global": + if self._adaptivity and not self._adaptivity_type == "global": raise Exception( "Load balancing can be done only with global adaptivity." ) except BaseException: self._logger.log_info_rank_zero( - "Micro Manager will not dynamically balance micro simulations based on the adaptivity computation." + "Micro Manager will not dynamically balance micro simulations." ) - if self._adaptivity_is_load_balancing: + if self._load_balancing: self._load_balancing_n = self._data["simulation_params"][ "load_balancing_settings" ]["every_n_time_windows"] @@ -417,29 +488,61 @@ def read_json_micro_manager(self): ) try: - self._load_balancing_threshold = self._data["simulation_params"][ + self._load_balancing_type = self._data["simulation_params"][ "load_balancing_settings" - ]["balancing_threshold"] - self._logger.log_info_rank_zero( - "Load balancing threshold: " + str(self._load_balancing_threshold) - ) + ]["type"] except BaseException: - self._logger.log_info_rank_zero( - "No load balancing threshold provided. The threshold will be set to 0." - ) + self._load_balancing_type = "time" + self._logger.log_info_rank_zero( + f"Load balancing will use {self._load_balancing_type} based balancing." + ) - try: - self._balance_inactive_sims = self._data["simulation_params"][ - "load_balancing_settings" - ]["balance_inactive_sims"] - if self._balance_inactive_sims: + if self._load_balancing_type == "active": + try: + self._load_balancing_threshold = self._data["simulation_params"][ + "load_balancing_settings" + ]["threshold"] + except BaseException: + self._load_balancing_threshold = 0 self._logger.log_info_rank_zero( - "Micro Manager will redistribute inactive simulations in the load balancing." + "Load balancing will use 0 threshold." ) - except BaseException: - self._logger.log_info_rank_zero( - "Micro Manager will not redistribute inactive simulations in the load balancing. Only active simulations will be redistributed. Note that this may significantly increase the communication cost of the adaptivity." - ) + + try: + self._load_balancing_balance_inactive_sims = self._data[ + "simulation_params" + ]["load_balancing_settings"]["balance_inactive_sims"] + except BaseException: + self._load_balancing_balance_inactive_sims = False + self._logger.log_info_rank_zero( + "Load balancing will not consider inactive simulations." + ) + else: + if ( + "threshold" + in self._data["simulation_params"]["load_balancing_settings"] + ): + self._logger.log_info_rank_zero( + 'Load balancing is not using active simulation balancing. Field "threshold" will be ignored.' + ) + if ( + "balance_inactive_sims" + in self._data["simulation_params"]["load_balancing_settings"] + ): + self._logger.log_info_rank_zero( + 'Load balancing is not using active simulation balancing. Field "balance_inactive_sims" will be ignored.' + ) + + if self._load_balancing_type == "time": + try: + self._load_balancing_partitioning = self._data["simulation_params"][ + "load_balancing_settings" + ]["partitioning"] + except BaseException: + self._logger.log_info_rank_zero( + "Partitioning type must be provided for time based load balancing. Defaulting to 'lpt'." + ) + self._load_balancing_partitioning = "lpt" try: if self._data["simulation_params"]["model_adaptivity"]: @@ -481,6 +584,28 @@ def read_json_micro_manager(self): "model_adaptivity_settings" ]["switching_function"] + if ( + "micro_stateless" + in self._data["simulation_params"]["model_adaptivity_settings"] + ): + self._m_adap_micro_stateless = self._data["simulation_params"][ + "model_adaptivity_settings" + ]["micro_stateless"] + else: + self._m_adap_micro_stateless = [False] * len( + self._m_adap_micro_file_names + ) + + for i in range(len(self._m_adap_micro_file_names)): + if self._m_adap_micro_stateless[i]: + self._logger.log_info_rank_zero( + f"Only creating one full instance of Micro Model {i}." + ) + else: + self._logger.log_info_rank_zero( + f"Creating full instance of Micro Model {i} per mesh vertex." + ) + if "interpolate_crash" in self._data["simulation_params"]: if self._data["simulation_params"]["interpolate_crash"]: self._interpolate_crash = True @@ -646,6 +771,28 @@ def get_ranks_per_axis(self): """ return self._ranks_per_axis + def get_decomposition_type(self): + """ + Get the type of domain decomposition. + + Returns + ------- + decomposition_type : str + Type of domain decomposition, can be either "uniform" or "non-uniform". + """ + return self._decomposition_type + + def get_minimum_access_region_size(self): + """ + Get the minimum access region size for non-uniform domain decomposition. + + Returns + ------- + minimum_access_region_size : list + List containing the minimum access region size in each axis for non-uniform domain decomposition. + """ + return self._minimum_access_region_size + def get_micro_file_name(self): """ Get the path to the Python script of the micro-simulation. @@ -657,6 +804,17 @@ def get_micro_file_name(self): """ return self._micro_file_name + def turn_on_micro_stateless(self): + """ + Boolean stating whether micro model is stateless or not. + + Returns + ------- + stateless : bool + True if micro model is stateless, False otherwise. + """ + return self._micro_stateless + def get_micro_output_n(self): """ Get the micro output frequency @@ -795,16 +953,16 @@ def is_adaptivity_required_in_every_implicit_iteration(self): """ return self._adaptivity_every_implicit_iteration - def is_adaptivity_with_load_balancing(self): + def turn_on_load_balancing(self): """ - Check if adaptivity computation needs to be done with load balancing. + Check if load balancing should be performed. Returns ------- - adaptivity_is_load_balancing : bool - True if adaptivity computation needs to be done with load balancing, False otherwise. + load_balancing : bool + True if load balancing needs to be done, False otherwise. """ - return self._adaptivity_is_load_balancing + return self._load_balancing def get_load_balancing_n(self): """ @@ -817,9 +975,20 @@ def get_load_balancing_n(self): """ return self._load_balancing_n + def get_load_balancing_type(self): + """ + Get load balancing type. + + Returns + ------- + type : str + Load balancing type. + """ + return self._load_balancing_type + def get_load_balancing_threshold(self): """ - Get the load balancing threshold to control how balanced the micro simulations need to be. + Get load balancing threshold. Returns ------- @@ -828,16 +997,27 @@ def get_load_balancing_threshold(self): """ return self._load_balancing_threshold - def balance_inactive_sims(self): + def turn_on_load_balancing_inactive(self): + """ + Check if load balancing should be performed on inactive micro simulations. + + Returns + ------- + balancing_inactive : bool + True if load balancing should consider inactive simulations, False otherwise. + """ + return self._load_balancing_balance_inactive_sims + + def get_load_balancing_partitioning(self): """ - Check if inactive simulations are to be redistributed in the load balancing. + Get the load balancing partitioning type Returns ------- - balance_inactive_sims : bool - True if inactive simulations are to be redistributed in the load balancing, False otherwise. + load_balancing_partitioning : str + Load balancing partitioning type """ - return self._balance_inactive_sims + return self._load_balancing_partitioning def initialize_sims_lazily(self): """ @@ -974,6 +1154,17 @@ def get_model_adaptivity_file_names(self): """ return self._m_adap_micro_file_names + def get_model_adaptivity_micro_stateless(self): + """ + List of boolean stating whether the respective micro model is stateless or not. + + Returns + ------- + stateless : list + True if micro model is stateless, False otherwise. + """ + return self._m_adap_micro_stateless + def get_model_adaptivity_switching_function(self): """ Get path to switching function file @@ -984,3 +1175,58 @@ def get_model_adaptivity_switching_function(self): String containing the path to the switching function file """ return self._m_adap_switching_function + + def get_tasking_num_workers(self): + """ + Get number of workers + + Returns + ------- + num_workers : int + Number of workers + """ + return self._task_num_workers + + def get_tasking_backend(self): + """ + Get backend type + + Returns + ------- + backend : str + either socket or mpi + """ + return self._task_backend + + def get_tasking_use_slurm(self): + """ + Get flag whether slurm is used + + Returns + ------- + use_slurm : bool + use slurm or not + """ + return self._task_is_slurm + + def get_tasking_hostfile(self): + """ + Get hostfile path for workers + + Returns + ------- + hostfile : str + Hostfile path for workers + """ + return self._task_pinning_hostfile + + def get_mpi_impl(self): + """ + Get mpi implementation type + + Returns + ------- + mpi_impl : str + mpi implementation type + """ + return self._task_mpi_impl diff --git a/micro_manager/domain_decomposition.py b/micro_manager/domain_decomposition.py index bbb97b16..e9bd7015 100644 --- a/micro_manager/domain_decomposition.py +++ b/micro_manager/domain_decomposition.py @@ -3,15 +3,20 @@ """ import numpy as np +from scipy.optimize import brentq +from micro_manager.config import Config +from typing import Callable class DomainDecomposer: - def __init__(self, rank, size) -> None: + def __init__(self, configurator: Config, rank: int, size: int) -> None: """ Class constructor. Parameters ---------- + configurator : object of class Config + Object which has getter functions to get parameters defined in the configuration file. rank : int MPI rank. size : int @@ -20,20 +25,41 @@ def __init__(self, rank, size) -> None: self._rank = rank self._size = size - def get_local_mesh_bounds(self, macro_bounds: list, ranks_per_axis: list) -> list: + self._ranks_per_axis = ( + configurator.get_ranks_per_axis() + ) # Check if ranks per axis is provided in the configuration file for parallel runs + + self._dims = len(self._ranks_per_axis) + + self._is_minimum_access_region_size_specified = False + + self._decomposition_type = configurator.get_decomposition_type() + + self._macro_bounds = configurator.get_macro_domain_bounds() + + self._get_local_mesh_bounds = self._get_local_mesh_bounds_variant() + + self._minimum_access_region_size: list = ( + configurator.get_minimum_access_region_size() + ) + if self._minimum_access_region_size: # if list is not empty + self._is_minimum_access_region_size_specified = True + + def get_local_mesh_bounds(self) -> list: """ - Decompose the macro domain equally among all ranks, if the Micro Manager is run in parallel. + Get the local mesh bounds for this rank based on the domain decomposition type specified in the configuration file. - Parameters - ---------- - macro_bounds : list - List containing upper and lower bounds of the macro domain. - Format in 2D is [x_min, x_max, y_min, y_max] - Format in 3D is [x_min, x_max, y_min, y_max, z_min, z_max] - ranks_per_axis : list - List containing axis wise ranks for a parallel run - Format in 2D is [ranks_x, ranks_y] - Format in 3D is [ranks_x, ranks_y, ranks_z] + Returns + ------- + mesh_bounds : list + List containing the upper and lower bounds of the domain pertaining to this rank. + Format is same as input parameter macro_bounds. + """ + return self._get_local_mesh_bounds() + + def _get_uniform_local_mesh_bounds(self) -> list: + """ + Decompose the macro domain equally among all ranks, if the Micro Manager is run in parallel. Returns ------- @@ -41,88 +67,189 @@ def get_local_mesh_bounds(self, macro_bounds: list, ranks_per_axis: list) -> lis List containing the upper and lower bounds of the domain pertaining to this rank. Format is same as input parameter macro_bounds. """ - if np.prod(ranks_per_axis) != self._size: + if np.prod(self._ranks_per_axis) != self._size: raise ValueError( "Total number of processors provided in the Micro Manager configuration and in the MPI execution command do not match." ) - dims = len(ranks_per_axis) - - if dims == 3: - for z in range(ranks_per_axis[2]): - for y in range(ranks_per_axis[1]): - for x in range(ranks_per_axis[0]): + if self._dims == 3: + for z in range(self._ranks_per_axis[2]): + for y in range(self._ranks_per_axis[1]): + for x in range(self._ranks_per_axis[0]): n = ( x - + y * ranks_per_axis[0] - + z * ranks_per_axis[0] * ranks_per_axis[1] + + y * self._ranks_per_axis[0] + + z * self._ranks_per_axis[0] * self._ranks_per_axis[1] ) if n == self._rank: rank_in_axis = [x, y, z] - elif dims == 2: - for y in range(ranks_per_axis[1]): - for x in range(ranks_per_axis[0]): - n = x + y * ranks_per_axis[0] + elif self._dims == 2: + for y in range(self._ranks_per_axis[1]): + for x in range(self._ranks_per_axis[0]): + n = x + y * self._ranks_per_axis[0] if n == self._rank: rank_in_axis = [x, y] else: raise ValueError("Domain decomposition only supports 2D and 3D cases.") - dx = [] - for d in range(dims): - dx.append( - abs(macro_bounds[d * 2 + 1] - macro_bounds[d * 2]) / ranks_per_axis[d] + mesh_bounds = [] + for d in range(self._dims): + dx = ( + abs(self._macro_bounds[d * 2 + 1] - self._macro_bounds[d * 2]) + / self._ranks_per_axis[d] ) - mesh_bounds = [] - for d in range(dims): if rank_in_axis[d] > 0: - mesh_bounds.append(macro_bounds[d * 2] + rank_in_axis[d] * dx[d]) - mesh_bounds.append(macro_bounds[d * 2] + (rank_in_axis[d] + 1) * dx[d]) + mesh_bounds.append(self._macro_bounds[d * 2] + rank_in_axis[d] * dx) + mesh_bounds.append( + self._macro_bounds[d * 2] + (rank_in_axis[d] + 1) * dx + ) elif rank_in_axis[d] == 0: - mesh_bounds.append(macro_bounds[d * 2]) - mesh_bounds.append(macro_bounds[d * 2] + dx[d]) + mesh_bounds.append(self._macro_bounds[d * 2]) + mesh_bounds.append(self._macro_bounds[d * 2] + dx) # Adjust the maximum bound to be exactly the domain size - if rank_in_axis[d] + 1 == ranks_per_axis[d]: - mesh_bounds[d * 2 + 1] = macro_bounds[d * 2 + 1] + if rank_in_axis[d] + 1 == self._ranks_per_axis[d]: + mesh_bounds[d * 2 + 1] = self._macro_bounds[d * 2 + 1] return mesh_bounds + def _get_nonuniform_local_mesh_bounds(self) -> list: + """ + Decompose the macro domain among all ranks with an non-uniform distribution, if the Micro Manager is run in parallel. + The non-uniform distribution is based on a geometric progression, where the size of the local mesh bounds increases + by a factor of 2 in each rank. This is just one of many possible non-uniform domain decompositions. + + Returns + ------- + mesh_bounds : list + List containing the upper and lower bounds of the domain pertaining to this rank. + Format is same as input parameter macro_bounds. + """ + if np.prod(self._ranks_per_axis) != self._size: + raise ValueError( + "Total number of processors provided in the Micro Manager configuration and in the MPI execution command do not match." + ) + + if self._dims == 3: + for z in range(self._ranks_per_axis[2]): + for y in range(self._ranks_per_axis[1]): + for x in range(self._ranks_per_axis[0]): + n = ( + x + + y * self._ranks_per_axis[0] + + z * self._ranks_per_axis[0] * self._ranks_per_axis[1] + ) + if n == self._rank: + rank_in_axis = [x, y, z] + elif self._dims == 2: + for y in range(self._ranks_per_axis[1]): + for x in range(self._ranks_per_axis[0]): + n = x + y * self._ranks_per_axis[0] + if n == self._rank: + rank_in_axis = [x, y] + else: + raise ValueError("Domain decomposition only supports 2D and 3D cases.") + + mesh_bounds = [] + multiplier = 2 # factor by which the local mesh bounds increase in each rank. 2 means geometric progression. + for d in range(self._dims): + macro_bounds_diff = abs( + self._macro_bounds[d * 2 + 1] - self._macro_bounds[d * 2] + ) + + dx0 = ( + macro_bounds_diff + * (multiplier - 1) + / (multiplier ** self._ranks_per_axis[d] - 1) + ) + + if self._is_minimum_access_region_size_specified: + if dx0 < self._minimum_access_region_size[d]: + dx0 = self._minimum_access_region_size[d] + n_ranks = self._ranks_per_axis[d] + + def _geom_sum_residual(r): + return dx0 * (r**n_ranks - 1) / (r - 1) - macro_bounds_diff + + # Find upper bracket where residual is positive + r_upper = 2.0 + while _geom_sum_residual(r_upper) <= 0: + r_upper *= 2.0 + + # When the minimum access region size is specified, + # the multiplier is numerically calculated such that the sum of + # the geometric progression of the local mesh bounds equals + # the macro domain size in that axis. + multiplier = brentq(_geom_sum_residual, 1.0 + 1e-12, r_upper) + + dx = np.zeros(self._ranks_per_axis[d]) + + for rank in range(self._ranks_per_axis[d]): + if rank == 0: + dx[rank] = dx0 + else: + dx[rank] = multiplier * dx[rank - 1] + + rank = rank_in_axis[d] + if rank == 0: + mesh_bounds.append(self._macro_bounds[d * 2]) + mesh_bounds.append(self._macro_bounds[d * 2] + dx[rank]) + + if rank > 0: + min_bound = self._macro_bounds[d * 2] + sum(dx[:rank]) + mesh_bounds.append(min_bound) + mesh_bounds.append(min_bound + dx[rank]) + + # Adjust the maximum bound of the access region of ranks on the boundary to be exactly the upper bound of the macro domain + if rank_in_axis[d] + 1 == self._ranks_per_axis[d]: + mesh_bounds[d * 2 + 1] = self._macro_bounds[d * 2 + 1] + + return mesh_bounds + + def _get_local_mesh_bounds_variant(self) -> Callable: + """ + Get uniform or nonuniform variant of calculating local mesh bounds + + Returns + ------- + get_local_mesh_bounds_variant : function + Function to calculate local mesh bounds based on the decomposition type specified in the configuration file. + """ + if self._decomposition_type == "uniform": + return self._get_uniform_local_mesh_bounds + elif self._decomposition_type == "nonuniform": + return self._get_nonuniform_local_mesh_bounds + else: + raise ValueError( + "Decomposition type can be either 'uniform' or 'nonuniform'." + ) + def get_local_sims_and_macro_coords( - self, macro_bounds: list, ranks_per_axis: list, macro_coords: np.ndarray - ) -> tuple[int, list]: + self, macro_coords: np.ndarray + ) -> tuple[int, list[np.ndarray]]: """ Decompose the micro simulations among all ranks based on their positions in the macro domain. - Parameters ---------- - macro_bounds : list - List containing upper and lower bounds of the macro domain. - Format in 2D is [x_min, x_max, y_min, y_max] - Format in 3D is [x_min, x_max, y_min, y_max, z_min, z_max] - ranks_per_axis : list - List containing axis wise ranks for a parallel run - Format in 2D is [ranks_x, ranks_y] - Format in 3D is [ranks_x, ranks_y, ranks_z] macro_coords : numpy.ndarray - The coordinates associated to the IDs and corresponding data values (dim * size) + Array containing the coordinates of the macro points corresponding to the micro simulations. Returns ------- micro_sims_on_rank : int - Number of micro simulations assigned to this rank. - macro_coords_on_this_rank : list - List of macro coordinates assigned to this rank. + Number of micro simulations pertaining to this rank. + macro_coords_on_this_rank : list of numpy.ndarray + List of coordinates of the macro points pertaining to this rank. """ - local_mesh_bounds = self.get_local_mesh_bounds(macro_bounds, ranks_per_axis) + local_mesh_bounds = self.get_local_mesh_bounds() macro_coords_on_this_rank = [] micro_sims_on_rank = 0 for position in macro_coords: inside = True - for d in range(len(ranks_per_axis)): + for d in range(self._dims): if not ( position[d] >= local_mesh_bounds[d * 2] and position[d] <= local_mesh_bounds[d * 2 + 1] @@ -134,3 +261,51 @@ def get_local_sims_and_macro_coords( micro_sims_on_rank += 1 return micro_sims_on_rank, macro_coords_on_this_rank + + def filter_duplicate_coords( + self, + all_coords: list, + all_ids: list, + ) -> tuple[np.ndarray, np.ndarray]: + """ + Filter out vertex coordinates that are already owned by a lower-ranked rank. + + When a macro-point lies exactly on the boundary between two rank bounding + boxes, preCICE returns it to both ranks. This function ensures every vertex + is processed by exactly one rank — the lowest-ranked rank that received it — + while preserving the preCICE ID-coord pairing. + + Parameters + ---------- + all_coords : list + List of numpy arrays, one per rank, containing vertex coordinates. + all_ids : list + List of arrays, one per rank, containing preCICE vertex IDs. + + Returns + ------- + filtered_coords : numpy.ndarray + Vertex coordinates with duplicates removed. + filtered_ids : numpy.ndarray + preCICE vertex IDs corresponding to the filtered coordinates. + """ + mesh_vertex_coords = np.array(all_coords[self._rank]) + mesh_vertex_ids = np.array(all_ids[self._rank]) + + seen_coords = set() + keep_mask = np.ones(len(mesh_vertex_coords), dtype=bool) + + for rank in range(self._size): + for i, coord in enumerate(all_coords[rank]): + coord_key = tuple(np.round(coord, decimals=10)) + if rank < self._rank: + # Mark coords already claimed by earlier ranks + seen_coords.add(coord_key) + elif rank == self._rank: + # Only keep coords not already claimed by earlier ranks + if coord_key in seen_coords: + keep_mask[i] = False + else: + seen_coords.add(coord_key) + + return mesh_vertex_coords[keep_mask], mesh_vertex_ids[keep_mask] diff --git a/micro_manager/interpolation.py b/micro_manager/interpolation.py index 966d961e..ee93a894 100644 --- a/micro_manager/interpolation.py +++ b/micro_manager/interpolation.py @@ -46,7 +46,7 @@ def get_nearest_neighbor_indices( return neighbor_indices def interpolate(self, neighbors: np.ndarray, point: np.ndarray, values): - """ + r""" Interpolate a value at a point using inverse distance weighting. (https://en.wikipedia.org/wiki/Inverse_distance_weighting) .. math:: f(x) = (\sum_{i=1}^{n} \frac{f_i}{\Vert x_i - x \Vert^2}) / (\sum_{j=1}^{n} \frac{1}{\Vert x_j - x \Vert^2}) diff --git a/micro_manager/load_balancing.py b/micro_manager/load_balancing.py new file mode 100644 index 00000000..cdd51a80 --- /dev/null +++ b/micro_manager/load_balancing.py @@ -0,0 +1,743 @@ +import time +from typing import Optional + +from mpi4py import MPI +import numpy as np + +from precice import Participant + +from micro_manager.config import Config +from micro_manager.adaptivity.adaptivity import AdaptivityCalculator +from micro_manager.model_manager import ModelManager +from micro_manager.tools.logging_wrapper import Logger +from micro_manager.tools.p2p import create_tag, get_ranks_of_sims + + +class LoadBalancer: + def __init__( + self, + precice_participant: Participant, + model_manager: ModelManager, + adaptivity_controller: Optional[AdaptivityCalculator], + state_loader: callable, + state_setter: callable, + log: Logger, + config: Config, + sim_list: list, + global_ids: list, + global_number_of_sims: int, + comm: MPI.Comm, + rank: int, + ): + """ + Constructs LoadBalancer. If load balancing is disabled, this will return a dummy instance + in which balancing request become NOOPs. + + Parameters + ---------- + precice_participant: Participant + preCICE participant object from coupling + model_manager: ModelManager + model manager object to construct instances + adaptivity_controller: Optional[AdaptivityCalculator] + handles adaptivity calculation if provided + state_loader: callable + loads state from micro simulation + state_setter: callable + sets state of micro simulation + log: Logger + logger object + config: Config + configuration object + sim_list: list + list of simulation objects + global_ids: list + list of global ids on this rank + global_number_of_sims: int + total number of simulations in this run + comm: MPI.Comm + used MPI communicator + rank: int + local rank + """ + self._enabled = config.turn_on_load_balancing() + self._precice_participant = precice_participant + self._model_manager = model_manager + self._adaptivity_controller = adaptivity_controller + self._state_loader = state_loader + self._state_setter = state_setter + self._log = log + self._config = config + self._sim_list = sim_list + self._global_ids = global_ids + self._global_number_of_sims = global_number_of_sims + self._comm = comm + self._rank = rank + + self._threshold = None # provided by sub-cls + self._balance_metric_local = dict() + self._balance_metric_global = np.zeros(global_number_of_sims) + self._partition_impl = self.get_partition_impl( + config.get_load_balancing_partitioning() + ) + + if ( + self._enabled + and adaptivity_controller is not None + and type(adaptivity_controller).__name__ != "GlobalAdaptivityCalculator" + ): + raise NotImplementedError( + "Adaptivity must be GlobalAdaptivity for Load Balancing" + ) + + def balance(self): + """ + Requests load balancing. If LoadBalancing is disabled, returns immediately. + """ + if not self._enabled: + return + + # self._precice_participant.start_profiling_section("micro_manager.solve.load_balancing.redistribute") + if np.allclose(self._balance_metric_global, 0): + self._balance_metric_global = self._balance_metric_global + 1 + self._redistribute() + # self._precice_participant.stop_last_profiling_section() + + def pre_sim_solve(self, gid: int): + """ + Notify load balancer that the micro simulation with the provided gid will start to run its solve method. + """ + self._balance_metric_local[gid] = time.time() + + def post_sim_solve(self, gid: int): + """ + Notify load balancer that the micro simulation with the provided gid has finished its solve method. + """ + self._balance_metric_local[gid] = time.time() - self._balance_metric_local[gid] + + def update(self): + """ + Needs to be called after all micro simulations have finished their solve method. + Updates the load balancing metric and shares it globally. + """ + + # used to distribute balancing metric + self._balance_metric_global[:] = 0 + tmp = self._comm.allgather(self._balance_metric_local) + for d in tmp: + for gid, val in d.items(): + self._balance_metric_global[gid] = val + + # clear local buffer + self._balance_metric_local.clear() + + # ================== + # PARTITIONING + # ================== + def get_partition_impl(self, name: str): + """ + Provides the selected partitioning algorithm. + + Parameters + ---------- + name: str + partitioning algorithm name + + Returns + ------- + function: callable + selected partitioning algorithm + """ + if name == "lpt": + return self.partition_lpt + else: + return self.partition_dummy + + def partition_lpt(self, n_parts: int, current_partitioning: np.ndarray): + """ + Partitions the recorded workload using the Longest-processing-time-first scheduling algorithm. + For more details see: https://en.wikipedia.org/wiki/Longest-processing-time-first_scheduling + + Parameters + ---------- + n_parts: int + number of partitions (number of ranks) + current_partitioning: np.ndarray + array of assignments of each work item to one partition < n_parts + + Returns + ------- + partitioning: np.ndarray + output of LPT algorithm + workload: np.ndarray + workload per partition + """ + sorted_workload_indices = np.argsort(self._balance_metric_global)[ + ::-1 + ] # descending + workload_per_partition = np.zeros(n_parts) + assignment = np.zeros(self._global_number_of_sims, dtype=np.int32) + + for idx in sorted_workload_indices: + # get current smallest partition + p = np.argmin(workload_per_partition) + # assign next largest work package + assignment[idx] = p + # update partition work load + workload_per_partition[p] += self._balance_metric_global[idx] + + return assignment, workload_per_partition + + def partition_dummy(self, n_parts: int, current_partitioning: np.ndarray): + """ + WARNING: Do not use this! This is only a dummy implementation that sends + the entire workload to the first partition. All others are empty. + + Parameters + ---------- + n_parts: int + number of partitions (number of ranks) + current_partitioning: np.ndarray + array of assignments of each work item to one partition < n_parts + + Returns + ------- + partition: np.ndarray + output of LPT algorithm + workload: np.ndarray + workload per partition + """ + # do not use this, just an example -> will send all to rank 0 + workload_per_partition = np.zeros(n_parts) + workload_per_partition[0] = np.sum(self._balance_metric_global) + return ( + np.zeros(self._global_number_of_sims, dtype=np.int32), + workload_per_partition, + ) + + # ================== + # HELPERS + # ================== + def _redistribute(self) -> None: + """ + Main implementation of load balancing. First computes the new partitioning. + Then send/receives micro simulations accordingly. + """ + self._precice_participant.start_profiling_section( + "micro_manager.solve.load_balancing.init" + ) + current_partitioning = get_ranks_of_sims( + self._global_ids, self._rank, self._comm, self._global_number_of_sims + ) + # self._precice_participant.start_profiling_section("micro_manager.solve.load_balancing.init.partition") + target_partitioning, work_loads = self._partition_impl( + self._comm.size, current_partitioning + ) + # self._precice_participant.stop_last_profiling_section() + send_map, recv_map = self._get_communication_maps( + current_partitioning, target_partitioning + ) + inactive_gids = self._get_global_inactive_gids() + self._precice_participant.stop_last_profiling_section() + + self._precice_participant.start_profiling_section( + "micro_manager.solve.load_balancing.comm" + ) + self._exchange_sims(send_map, recv_map, {gid: True for gid in inactive_gids}) + self._precice_participant.stop_last_profiling_section() + + sims_per_rank = self._comm.gather(len(self._sim_list), 0) + self._log.log_info_rank_zero( + f"Load Balancing Number of Simulations per Rank: {sims_per_rank}" + ) + + def _get_communication_maps( + self, current_partitioning: np.ndarray, target_partitioning: np.ndarray + ) -> tuple: + """ + Create dictionaries which map global IDs of simulations to ranks for sending and receiving. + + Parameters + ---------- + current_partitioning : np.ndarray + Current assignment of simulations + target_partitioning : np.ndarray + Target assignment of simulations + + Returns + ------- + tuple of dicts + send_map : dict + keys are global IDs, values are target ranks + recv_map : dict + keys are global IDs, values are source ranks + """ + send_map = dict() + recv_map = dict() + + for gid, target_rank in enumerate(target_partitioning): + if current_partitioning[gid] == target_rank: + continue + if current_partitioning[gid] == self._rank: + send_map[gid] = target_rank + continue + if target_rank == self._rank: + recv_map[gid] = current_partitioning[gid] + + return send_map, recv_map + + def _get_global_active_gids(self): + """ + Get global IDs of all active gids. This is based on local ids. + + Returns + ------- + active_gids: list[int] + list of global active gids + """ + # local count + active_gid_arr = None + if self._adaptivity_controller is not None: + active_gid_arr = [ + self._global_ids[i] + for i in self._adaptivity_controller.get_active_sim_local_ids() + ] + else: + active_gid_arr = self._global_ids + + # bcast to all and merge + tmp = self._comm.allgather(active_gid_arr) + global_active_gids = [] + for l in tmp: + global_active_gids.extend(l) + return global_active_gids + + def _get_global_inactive_gids(self): + """ + Get global IDs of all inactive gids. This is based on local ids. + + Returns + ------- + inactive_gids: list[int] + list of global inactive gids + """ + global_active_gids = set(self._get_global_active_gids()) + global_inactive_gids = set(np.arange(self._global_number_of_sims)).difference( + global_active_gids + ) + return list(global_inactive_gids) + + def _exchange_sims(self, send_map, recv_map, inactive_map={}): + """ + Move active micro simulations between ranks. + Sends state+gid if simulation is active, None+gid otherwise. + + Parameters + ---------- + send_map : dict + keys are global IDs of sim states to send, values are ranks to send the sims to + recv_map : dict + keys are global IDs of sim states to receive, values are ranks to receive from + inactive_map : dict + keys are global IDs of inactive sim states, values are bool + """ + # Asynchronous send operations + send_reqs = [] + for gid, send_rank in send_map.items(): + tag = create_tag(gid, self._rank, send_rank) + lid = self._global_ids.index(gid) + + # prepare send data + is_inactive = inactive_map[gid] if gid in inactive_map else False + cls_name = None if is_inactive else self._sim_list[lid].name + is_stateless = ( + None if is_inactive else self._model_manager.is_stateless(cls_name) + ) + send_data = ( + is_inactive, + is_stateless, + None + if is_stateless or is_inactive + else self._state_loader(self._sim_list[lid]), + cls_name, + gid, + ) + + req = self._comm.isend(send_data, dest=send_rank, tag=tag) + send_reqs.append(req) + + # Asynchronous receive operations + recv_reqs = [] + for gid, recv_rank in recv_map.items(): + tag = create_tag(gid, recv_rank, self._rank) + bufsize = ( + 1 << 30 + ) # allocate and use a temporary 1 MiB buffer size https://github.com/mpi4py/mpi4py/issues/389 + req = self._comm.irecv(bufsize, source=recv_rank, tag=tag) + recv_reqs.append(req) + + # Wait for all non-blocking communication to complete + MPI.Request.Waitall(send_reqs) + + # Delete the simulations which no longer exist on this rank + for gid in send_map.keys(): + lid = self._global_ids.index(gid) + is_active = gid not in inactive_map + if is_active: + self._sim_list[lid].destroy() + del self._sim_list[lid] + self._global_ids.remove(gid) + if self._adaptivity_controller is not None: + self._adaptivity_controller.set_is_on_rank(gid, False) + + # Create simulations and set them to the received states + for req in recv_reqs: + is_inactive, is_stateless, state, cls_name, gid = req.wait() + self._sim_list.append( + None + if is_inactive + else self._model_manager.get_instance_by_name(gid, cls_name) + ) + if not is_stateless and state is not None: + self._state_setter(self._sim_list[-1], state) + self._global_ids.append(gid) + if self._adaptivity_controller is not None: + self._adaptivity_controller.set_is_on_rank(gid, True) + + +class ActiveBalancer(LoadBalancer): + """ + ActiveBalancer will attempt to balance the number of active micro simulations between ranks. + """ + + def __init__( + self, + precice_participant: Participant, + model_manager: ModelManager, + adaptivity_controller: Optional[AdaptivityCalculator], + state_loader: callable, + state_setter: callable, + log: Logger, + config: Config, + sim_list: list, + global_ids: list, + global_number_of_sims: int, + comm: MPI.Comm, + rank: int, + ): + super().__init__( + precice_participant, + model_manager, + adaptivity_controller, + state_loader, + state_setter, + log, + config, + sim_list, + global_ids, + global_number_of_sims, + comm, + rank, + ) + self._partition_impl = lambda a, b: (None, None) + self._threshold = config.get_load_balancing_threshold() + self._balance_inactive_sims = config.turn_on_load_balancing_inactive() + self._bypass_skip = False # used for testing + self._bypass_active = False # used for testing + + if adaptivity_controller is None: + raise ValueError( + "Active Count balancing requires GlobalAdaptivityCalculator" + ) + + def pre_sim_solve(self, gid): + pass + + def post_sim_solve(self, gid): + pass + + def update(self): + pass + + def _get_active_exchange_counts(self): + avg_active_sims = ( + np.count_nonzero(self._adaptivity_controller._is_sim_active) + / self._comm.size + ) + f_avg_active_sims = np.floor(avg_active_sims - self._threshold) + c_avg_active_sims = np.ceil(avg_active_sims + self._threshold) + + active_sims_local_ids = self._adaptivity_controller.get_active_sim_local_ids() + n_active_sims_local = len(active_sims_local_ids) + send_sims = 0 # Sims that this rank wants to send + recv_sims = 0 # Sims that this rank wants to receive + + if f_avg_active_sims == c_avg_active_sims: + if n_active_sims_local < avg_active_sims: + recv_sims = int(avg_active_sims) - n_active_sims_local + elif n_active_sims_local > avg_active_sims: + send_sims = n_active_sims_local - int(avg_active_sims) + else: + if n_active_sims_local < f_avg_active_sims: + recv_sims = f_avg_active_sims - n_active_sims_local + elif n_active_sims_local == f_avg_active_sims: + recv_sims += 1 + elif n_active_sims_local > c_avg_active_sims: + send_sims = n_active_sims_local - c_avg_active_sims + elif n_active_sims_local == c_avg_active_sims: + send_sims += 1 + + # Number of active sims that each rank wants to send and receive + global_send_sims = self._comm.allgather(send_sims) + global_recv_sims = self._comm.allgather(recv_sims) + + n_global_send_sims = sum(global_send_sims) + n_global_recv_sims = sum(global_recv_sims) + + return ( + global_send_sims, + global_recv_sims, + n_global_send_sims, + n_global_recv_sims, + ) + + def _get_active_comm_maps(self, global_send_sims: list, global_recv_sims: list): + """ + Create dictionaries which map global IDs of simulations to ranks for sending and receiving. + + Parameters + ---------- + global_send_sims : list + Number of simulations that each rank sends. + global_recv_sims : list + Number of simulations that each rank receives. + + Returns + ------- + tuple of dicts + send_map : dict + keys are global IDs of sim states to send, values are ranks to send the sims to + recv_map : dict + keys are global IDs of sim states to receive, values are ranks to receive from + """ + active_sims_global_ids = list( + self._adaptivity_controller.get_active_sim_global_ids() + ) + rank_wise_global_ids_of_active_sims = self._comm.allgather( + active_sims_global_ids + ) + + # Keys are ranks sending sims, values are lists of tuples: (list of global IDs to send, the rank to send them to) + global_send_map: dict[int, list] = dict() + + # Keys are ranks receiving sims, values are lists of tuples: (list of global IDs to receive, the rank to receive them from) + global_recv_map: dict[int, list] = dict() + + for rank in [i for i, e in enumerate(global_send_sims) if e != 0]: + global_send_map[rank] = [] + + for rank in [i for i, e in enumerate(global_recv_sims) if e != 0]: + global_recv_map[rank] = [] + + send_ranks = list(global_send_map.keys()) + recv_ranks = list(global_recv_map.keys()) + + count = 0 + recv_rank = recv_ranks[count] + + for send_rank in send_ranks: + sims = global_send_sims[send_rank] + while sims > 0: + if global_recv_sims[recv_rank] <= sims: + # Get the global IDs to move + global_ids_of_sims_to_move = rank_wise_global_ids_of_active_sims[ + send_rank + ][0 : int(global_recv_sims[recv_rank])] + + global_send_map[send_rank].append( + (global_ids_of_sims_to_move, recv_rank) + ) + + global_recv_map[recv_rank].append( + (global_ids_of_sims_to_move, send_rank) + ) + + sims -= global_recv_sims[recv_rank] + + # Remove the global IDs which are already mapped for moving + del rank_wise_global_ids_of_active_sims[send_rank][ + 0 : int(global_recv_sims[recv_rank]) + ] + + if count < len(recv_ranks) - 1: + count += 1 + recv_rank = recv_ranks[count] + + elif global_recv_sims[recv_rank] > sims: + # Get the global IDs to move + global_ids_of_sims_to_move = rank_wise_global_ids_of_active_sims[ + send_rank + ][0 : int(sims)] + + global_send_map[send_rank].append( + (global_ids_of_sims_to_move, recv_rank) + ) + + global_recv_map[recv_rank].append( + (global_ids_of_sims_to_move, send_rank) + ) + + global_recv_sims[recv_rank] -= sims + + # Remove the global IDs which are already mapped for moving + del rank_wise_global_ids_of_active_sims[send_rank][0 : int(sims)] + + sims = 0 + + # keys are global IDs of sim states to send, values are ranks to send the sims to + send_map: dict[int, int] = dict() + + # keys are global IDs of sim states to receive, values are ranks to receive from + recv_map: dict[int, int] = dict() + + if self._rank in global_send_map: + for send_info in global_send_map[self._rank]: + send_rank = send_info[1] + for gid in send_info[0]: + send_map[gid] = send_rank + + if self._rank in global_recv_map: + for recv_info in global_recv_map[self._rank]: + recv_rank = recv_info[1] + for gid in recv_info[0]: + recv_map[gid] = recv_rank + + return send_map, recv_map + + def _get_inactive_comm_maps(self): + send_map: dict[int, int] = dict() + recv_map: dict[int, int] = dict() + ranks_of_sims = get_ranks_of_sims( + self._global_ids, self._rank, self._comm, self._global_number_of_sims + ) + global_ids_of_inactive_sims = self._get_global_inactive_gids() + + for gid in global_ids_of_inactive_sims: + assoc_active_gid = self._adaptivity_controller._sim_is_associated_to[gid] + rank_of_inactive_sim = ranks_of_sims[gid] + rank_of_assoc_active_sim = ranks_of_sims[assoc_active_gid] + if rank_of_inactive_sim != rank_of_assoc_active_sim: + if rank_of_inactive_sim == self._rank: + send_map[gid] = rank_of_assoc_active_sim + if rank_of_assoc_active_sim == self._rank: + recv_map[gid] = rank_of_inactive_sim + + return send_map, recv_map + + @staticmethod + def _correct_active_exchange_data( + global_send_sims, global_recv_sims, n_global_send_sims, n_global_recv_sims + ): + if n_global_send_sims < n_global_recv_sims: + excess_recv_sims = n_global_recv_sims - n_global_send_sims + while excess_recv_sims > 0: + for i, e in enumerate(global_recv_sims): + if e <= 0: + continue + # Remove the excess receive request from the rank + global_recv_sims[i] -= 1 + excess_recv_sims -= 1 + if excess_recv_sims == 0: + break + elif n_global_send_sims > n_global_recv_sims: + excess_send_sims = n_global_send_sims - n_global_recv_sims + while excess_send_sims > 0: + for i, e in enumerate(global_send_sims): + if e <= 0: + continue + # Remove the excess send request + global_send_sims[i] -= 1 + excess_send_sims -= 1 + if excess_send_sims == 0: + break + + def _get_communication_maps(self, *args, **kwargs): + send_map: dict[int, int] = dict() + recv_map: dict[int, int] = dict() + + if not self._bypass_active: + ( + global_send_sims, + global_recv_sims, + n_global_send_sims, + n_global_recv_sims, + ) = self._get_active_exchange_counts() + + if ( + n_global_send_sims == 0 + and n_global_recv_sims == 0 + and not self._bypass_skip + ): + self._log.log_warning_rank_zero( + "It appears that the micro simulations are already fairly balanced. No load balancing will be done. Try changing the threshold value to induce load balancing." + ) + return send_map, recv_map + if n_global_send_sims != 0 or n_global_recv_sims != 0: + ActiveBalancer._correct_active_exchange_data( + global_send_sims, + global_recv_sims, + n_global_send_sims, + n_global_recv_sims, + ) + send_map_active, recv_map_active = self._get_active_comm_maps( + global_send_sims, global_recv_sims + ) + send_map.update(send_map_active) + recv_map.update(recv_map_active) + + # if requested, also balance inactive simulations if there was a change in active simulations + if self._balance_inactive_sims: + send_map_inactive, recv_map_inactive = self._get_inactive_comm_maps() + send_map.update(send_map_inactive) + recv_map.update(recv_map_inactive) + + return send_map, recv_map + + +def create_load_balancer( + precice_participant: Participant, + model_manager: ModelManager, + adaptivity_controller: Optional[AdaptivityCalculator], + state_loader: callable, + state_setter: callable, + log: Logger, + config: Config, + sim_list: list, + global_ids: list, + global_number_of_sims: int, + comm: MPI.Comm, + rank: int, +) -> LoadBalancer: + lb_type = config.get_load_balancing_type() + + if lb_type == "time": + lb_cls = LoadBalancer + elif lb_type == "active": + lb_cls = ActiveBalancer + else: + raise RuntimeError(f"Unknown load balancing type: {lb_type}") + + return lb_cls( + precice_participant, + model_manager, + adaptivity_controller, + state_loader, + state_setter, + log, + config, + sim_list, + global_ids, + global_number_of_sims, + comm, + rank, + ) diff --git a/micro_manager/micro_manager.py b/micro_manager/micro_manager.py index 6f924b02..64239a35 100644 --- a/micro_manager/micro_manager.py +++ b/micro_manager/micro_manager.py @@ -12,10 +12,8 @@ Detailed documentation: https://precice.org/tooling-micro-manager-overview.html """ -import importlib import os import sys -import inspect from typing import Callable import numpy as np from psutil import Process @@ -25,18 +23,17 @@ import precice +from .model_manager import ModelManager from .micro_manager_base import MicroManager -from .adaptivity.global_adaptivity import GlobalAdaptivityCalculator -from .adaptivity.local_adaptivity import LocalAdaptivityCalculator -from .adaptivity.global_adaptivity_lb import GlobalAdaptivityLBCalculator from .adaptivity.model_adaptivity import ModelAdaptivity +from .adaptivity.adaptivity_selection import create_adaptivity_calculator from .domain_decomposition import DomainDecomposer - -from .micro_simulation import create_simulation_class +from .tasking.connection import spawn_local_workers +from .micro_simulation import create_simulation_class, load_backend_class from .tools.logging_wrapper import Logger - +from .load_balancing import create_load_balancer try: from .interpolation import Interpolation @@ -86,11 +83,6 @@ def __init__(self, config_file: str, log_file: str = "") -> None: self._macro_mesh_name = self._config.get_macro_mesh_name() - self._macro_bounds = self._config.get_macro_domain_bounds() - - if self._is_parallel: # Simulation is run in parallel - self._ranks_per_axis = self._config.get_ranks_per_axis() - # Parameter for interpolation in case of a simulation crash self._interpolate_crashed_sims = self._config.interpolate_crashed_micro_sim() if self._interpolate_crashed_sims: @@ -110,10 +102,6 @@ def __init__(self, config_file: str, log_file: str = "") -> None: self._is_adaptivity_on = self._config.turn_on_adaptivity() - self._is_adaptivity_with_load_balancing = ( - self._config.is_adaptivity_with_load_balancing() - ) - if self._is_adaptivity_on: self._data_for_adaptivity: dict[str, list] = dict() @@ -134,9 +122,6 @@ def __init__(self, config_file: str, log_file: str = "") -> None: self._config.is_adaptivity_required_in_every_implicit_iteration() ) - if self._is_adaptivity_with_load_balancing: - self._load_balancing_n = self._config.get_load_balancing_n() - self._adaptivity_n = self._config.get_adaptivity_n() self._adaptivity_output_type = self._config.get_adaptivity_output_type() @@ -156,6 +141,19 @@ def __init__(self, config_file: str, log_file: str = "") -> None: self._t = 0 # global time self._n = 0 # sim-step + self._model_manager = ModelManager() + self._conn = None + self.state_loader = lambda sim: sim.get_state() + self.state_setter = lambda sim, state: sim.set_state(state) + if self._is_model_adaptivity_on: + self.state_loader = lambda sim: sim.attachments + self.state_setter = lambda sim, state: sim.attachments.update(state) + self._is_load_balancing = ( + self._config.turn_on_load_balancing() and self._is_parallel + ) + self._load_balancing_n = self._config.get_load_balancing_n() + self.load_balancing = None + # ************** # Public methods # ************** @@ -182,6 +180,7 @@ def solve(self) -> None: dt = min(self._participant.get_max_time_step_size(), self._micro_dt) first_iteration = True + lb_counter = -1 if self._is_adaptivity_on: # Log initial adaptivity metrics @@ -204,50 +203,60 @@ def solve(self) -> None: self._micro_sims, self._data_for_adaptivity, ) - - # Write a checkpoint if a simulation is just activated. - # This checkpoint will be asynchronous to the checkpoints written at the start of the time window. - for i in range(self._local_number_of_sims): - if sim_states_cp[i] is None and self._micro_sims[i]: - sim_states_cp[i] = self._micro_sims[i].get_state() - active_sim_gids = ( self._adaptivity_controller.get_active_sim_global_ids() ) - for gid in active_sim_gids: self._micro_sims_active_steps[gid] += 1 - self._participant.stop_last_profiling_section() - - if self._is_adaptivity_with_load_balancing: - if self._n % self._load_balancing_n == 0 and first_iteration: - self._participant.start_profiling_section( - "micro_manager.solve.load_balancing" - ) - - self._adaptivity_controller.redistribute_sims(self._micro_sims) - - self._local_number_of_sims = len(self._global_ids_of_local_sims) + # Write a checkpoint if a simulation is just activated. + # This checkpoint will be asynchronous to the checkpoints written at the start of the time window. + if self._is_model_adaptivity_on: + self._model_adaptivity_controller.update_states( + self._micro_sims, active_sim_gids + ) + for i in range(self._local_number_of_sims): + if sim_states_cp[i] is None and self._micro_sims[i]: + sim_states_cp[i] = self.state_loader(self._micro_sims[i]) - # Reset simulation state checkpoints after load balancing - sim_states_cp = [None] * self._local_number_of_sims + self._participant.stop_last_profiling_section() + # handle load balancing, in first iteration all sims are assumed to have same cost + performed_lb = False + if lb_counter % self._load_balancing_n == 0 or first_iteration: + # self._participant.start_profiling_section("micro_manager.solve.load_balancing") + self.load_balancing.balance() + self._local_number_of_sims = len(self._global_ids_of_local_sims) + self._is_rank_empty = self._local_number_of_sims == 0 + # Reset simulation state checkpoints after load balancing + sim_states_cp = [None] * self._local_number_of_sims + if self._is_adaptivity_on: for name in self._adaptivity_data_names: self._data_for_adaptivity[name] = [ 0 ] * self._local_number_of_sims - - # Reset simulation crash state information after load balancing - self._has_sim_crashed = [False] * self._local_number_of_sims - - self._participant.stop_last_profiling_section() + # Reset simulation crash state information after load balancing + self._has_sim_crashed = [False] * self._local_number_of_sims + # self._participant.stop_last_profiling_section() + performed_lb = True + lb_counter += 1 # Write a checkpoint - if self._participant.requires_writing_checkpoint(): + if self._participant.requires_writing_checkpoint() or performed_lb: + if self._is_model_adaptivity_on: + active_sim_gids = None + if self._is_adaptivity_on: + active_sim_gids = ( + self._adaptivity_controller.get_active_sim_local_ids() + ) + self._model_adaptivity_controller.update_states( + self._micro_sims, active_sim_gids + ) for i in range(self._local_number_of_sims): sim_states_cp[i] = ( - self._micro_sims[i].get_state() if self._micro_sims[i] else None + self.state_loader(self._micro_sims[i]) + if self._micro_sims[i] + else None ) micro_sims_input = self._read_data_from_precice(dt) @@ -257,12 +266,13 @@ def solve(self) -> None: ) micro_sims_output = micro_sim_solve(micro_sims_input, dt) + self.load_balancing.update() self._participant.stop_last_profiling_section() - if self._is_adaptivity_with_load_balancing: + if self._is_load_balancing: for i in range(self._local_number_of_sims): - micro_sims_output[i]["Rank"] = self._rank + micro_sims_output[i]["rank_of_sim"] = self._rank # Check if more than a certain percentage of the micro simulations have crashed and terminate if threshold is exceeded if self._interpolate_crashed_sims: @@ -295,12 +305,22 @@ def solve(self) -> None: if self._participant.requires_reading_checkpoint(): for i in range(self._local_number_of_sims): if self._micro_sims[i]: - self._micro_sims[i].set_state(sim_states_cp[i]) + self.state_setter(self._micro_sims[i], sim_states_cp[i]) + + if self._is_model_adaptivity_on: + active_sim_gids = None + if self._is_adaptivity_on: + active_sim_gids = ( + self._adaptivity_controller.get_active_sim_local_ids() + ) + self._model_adaptivity_controller.write_back_states( + self._micro_sims, active_sim_gids + ) + first_iteration = False - if ( - self._participant.is_time_window_complete() - ): # Time window has converged, now micro output can be generated + # Time window has converged, now micro output can be generated + if self._participant.is_time_window_complete(): self._t += dt # Update time to the end of the time window self._n += 1 # Update time step to the end of the time window @@ -391,6 +411,8 @@ def solve(self) -> None: for i, rss_mb in enumerate(avg_mem_usage): writer.writerow([mem_usage_n[i], rss_mb]) + if self._conn is not None: + self._conn.close() self._participant.finalize() def initialize(self) -> None: @@ -407,31 +429,26 @@ def initialize(self) -> None: ) # Decompose the macro-domain and set the mesh access region for each partition in preCICE - if not len(self._macro_bounds) / 2 == self._participant.get_mesh_dimensions( - self._macro_mesh_name - ): + if not len( + self._config.get_macro_domain_bounds() + ) / 2 == self._participant.get_mesh_dimensions(self._macro_mesh_name): raise Exception("Provided macro mesh bounds are of incorrect dimension") if self._is_parallel: - if not len(self._ranks_per_axis) == self._participant.get_mesh_dimensions( - self._macro_mesh_name - ): + if not len( + self._config.get_ranks_per_axis() + ) == self._participant.get_mesh_dimensions(self._macro_mesh_name): raise Exception( "Provided ranks combination is of incorrect dimension" " and does not match the dimensions of the macro mesh." ) - domain_decomposer = DomainDecomposer( - self._rank, - self._size, - ) + domain_decomposer = DomainDecomposer(self._config, self._rank, self._size) - if self._is_parallel and not self._is_adaptivity_with_load_balancing: - coupling_mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds, self._ranks_per_axis - ) + if self._is_parallel and not self._is_load_balancing: + coupling_mesh_bounds = domain_decomposer.get_local_mesh_bounds() else: # When serial or load balancing, the whole macro domain is assigned to one/each rank - coupling_mesh_bounds = self._macro_bounds + coupling_mesh_bounds = self._config.get_macro_domain_bounds() self._participant.set_mesh_access_region( self._macro_mesh_name, coupling_mesh_bounds @@ -451,30 +468,43 @@ def initialize(self) -> None: self._mesh_vertex_coords, ) = self._participant.get_mesh_vertex_ids_and_coordinates(self._macro_mesh_name) + if self._is_parallel and not self._is_load_balancing: + # Gather all vertex coords and IDs from all ranks onto all ranks, + # filter out coords already claimed by lower-ranked ranks. + # When load balancing, all ranks receive all coords. No duplicates can arise. + all_coords = self._comm.allgather(self._mesh_vertex_coords) + all_ids = self._comm.allgather(self._mesh_vertex_ids) + + ( + self._mesh_vertex_coords, + self._mesh_vertex_ids, + ) = domain_decomposer.filter_duplicate_coords(all_coords, all_ids) + if self._mesh_vertex_coords.size == 0: - raise Exception("Macro mesh has no vertices.") + if self._is_parallel: + self._is_rank_empty = True + self._logger.log_warning( + "The access region of this rank has no macro-scale vertices. This rank will not have any micro simulations. To avoid this, change the domain decomposition" + ) + if self._lazy_init: + raise Exception( + "The macro mesh has no vertices in the specified access region, but lazy initialization is turned on. Lazy initialization cannot be used if there are no vertices in the access region, as there would be no data to compute the adaptivity and determine which simulations to initialize." + ) + else: + raise Exception( + "The macro mesh has no vertices in the specified access region." + ) - if self._is_adaptivity_with_load_balancing: + if self._is_load_balancing and self._is_parallel: ( self._local_number_of_sims, local_macro_coords, ) = domain_decomposer.get_local_sims_and_macro_coords( - self._macro_bounds, self._ranks_per_axis, self._mesh_vertex_coords + self._mesh_vertex_coords ) else: self._local_number_of_sims, _ = self._mesh_vertex_coords.shape - if self._local_number_of_sims == 0: - if self._is_parallel: - self._logger.log_info( - "Rank {} has no micro simulations and hence will not do any computation.".format( - self._rank - ) - ) - self._is_rank_empty = True - else: - raise Exception("Micro Manager has no micro simulations.") - nms_all_ranks = np.zeros(self._size, dtype=np.int64) # Gather number of micro simulations that each rank has, because this rank needs to know how many micro # simulations have been created by previous ranks, so that it can set @@ -518,7 +548,7 @@ def initialize(self) -> None: # Create lists of global IDs self._global_ids_of_local_sims = [] # DECLARATION - if self._is_adaptivity_with_load_balancing: + if self._is_load_balancing: # Create a set of global coordinate indices for faster lookup coord_to_index = { tuple(coord): i for i, coord in enumerate(self._mesh_vertex_coords) @@ -539,72 +569,69 @@ def initialize(self) -> None: if self._interpolate_crashed_sims: self._interpolant = Interpolation(self._logger) + # Setup remote workers + base_dir = os.path.dirname(os.path.abspath(__file__)) + worker_exec = os.path.join(base_dir, "tasking", "worker_main.py") + num_ranks = self._config.get_tasking_num_workers() + self._conn = spawn_local_workers( + worker_exec, + num_ranks, + self._config.get_tasking_backend(), + self._config.get_tasking_use_slurm(), + self._config.get_mpi_impl(), + self._config.get_tasking_hostfile(), + ) + + # load micro sim micro_problem_cls = None if self._is_model_adaptivity_on: self._model_adaptivity_controller: ModelAdaptivity = ModelAdaptivity( - self._config, self._rank, self._log_file + self._model_manager, + self._config, + self._comm, + self._rank, + self._log_file, + self._conn, + num_ranks, ) micro_problem_cls = ( self._model_adaptivity_controller.get_resolution_sim_class(0) ) else: - micro_problem_base = getattr( - importlib.import_module( - self._config.get_micro_file_name(), "MicroSimulation" - ), - "MicroSimulation", - ) + micro_problem_base = load_backend_class(self._config.get_micro_file_name()) micro_problem_cls = create_simulation_class( - micro_problem_base, "MicroSimulationDefault" + self._logger, + micro_problem_base, + self._config.get_micro_file_name(), + self._config.get_tasking_num_workers(), + self._conn, + "MicroSimulationDefault", + ) + self._model_manager.register( + micro_problem_cls, self._config.turn_on_micro_stateless() ) # Create micro simulation objects - self._micro_sims = [0] * self._local_number_of_sims + self._micro_sims = [None] * self._local_number_of_sims if not self._lazy_init: for i in range(self._local_number_of_sims): - self._micro_sims[i] = micro_problem_cls( - self._global_ids_of_local_sims[i] + self._micro_sims[i] = self._model_manager.get_instance( + self._global_ids_of_local_sims[i], micro_problem_cls ) if self._is_adaptivity_on: - if self._config.get_adaptivity_type() == "local": - self._adaptivity_controller: LocalAdaptivityCalculator = ( - LocalAdaptivityCalculator( - self._config, - self._local_number_of_sims, - self._logger, - self._rank, - self._comm, - micro_problem_cls, - ) - ) - elif self._config.get_adaptivity_type() == "global": - if self._is_adaptivity_with_load_balancing: - self._adaptivity_controller: GlobalAdaptivityLBCalculator = ( - GlobalAdaptivityLBCalculator( - self._config, - self._global_number_of_sims, - self._global_ids_of_local_sims, - self._participant, - self._logger, - self._rank, - self._comm, - micro_problem_cls, - ) - ) - else: - self._adaptivity_controller: GlobalAdaptivityCalculator = ( - GlobalAdaptivityCalculator( - self._config, - self._global_number_of_sims, - self._global_ids_of_local_sims, - self._participant, - self._logger, - self._rank, - self._comm, - micro_problem_cls, - ) - ) + self._adaptivity_controller = create_adaptivity_calculator( + self._config, + self._local_number_of_sims, + self._global_number_of_sims, + self._global_ids_of_local_sims, + self._participant, + self._logger, + self._rank, + self._comm, + micro_problem_cls, + self._model_manager, + ) self._micro_sims_active_steps = np.zeros( self._global_number_of_sims @@ -613,23 +640,31 @@ def initialize(self) -> None: self._micro_sims_init = False # DECLARATION # Read initial data from preCICE, if it is available - initial_data = self._read_data_from_precice(dt=0) + initial_macro_data = self._read_data_from_precice(dt=0) first_id = 0 # 0 if lazy initialization is off, otherwise the first active simulation ID micro_sims_to_init = range( 1, self._local_number_of_sims ) # All sims if lazy init is off, otherwise all active simulations - if not initial_data: + # Additional bool to check if there are sims to init + are_there_sims_to_init = True + + if not initial_macro_data: is_initial_data_available = False + + if self._lazy_init: + raise Exception( + "Initial macro data is required for lazy initialization." + ) else: is_initial_data_available = True - if ( - self._lazy_init - ): # For lazy initialization, compute adaptivity with the initial macro data + + # For lazy initialization, compute adaptivity with the initial macro data + if self._lazy_init: for i in range(self._local_number_of_sims): for name in self._adaptivity_macro_data_names: - self._data_for_adaptivity[name][i] = initial_data[i][name] + self._data_for_adaptivity[name][i] = initial_macro_data[i][name] self._adaptivity_controller.compute_adaptivity( self._micro_dt, self._micro_sims, self._data_for_adaptivity @@ -641,78 +676,58 @@ def initialize(self) -> None: self._logger.log_info( "There are no active simulations on this rank." ) - return - - for i in active_sim_lids: - self._micro_sims[i] = micro_problem_cls( - self._global_ids_of_local_sims[i] - ) - - first_id = active_sim_lids[0] # First active simulation ID - micro_sims_to_init = ( - active_sim_lids # Only active simulations will be initialized - ) - - # Boolean which states if the initialize() method of the micro simulation requires initial data - sim_requires_init_data = False + micro_sims_to_init = [] + are_there_sims_to_init = False + else: + for i in active_sim_lids: + self._micro_sims[i] = micro_problem_cls( + self._global_ids_of_local_sims[i] + ) - # Check if provided micro simulation has an initialize() method - if hasattr(micro_problem_cls, "initialize") and callable( - getattr(micro_problem_cls, "initialize") - ): - self._micro_sims_init = True # Starting value before setting + for i in active_sim_lids: + self._micro_sims[i] = self._model_manager.get_instance( + self._global_ids_of_local_sims[i], micro_problem_cls + ) - try: # Try to get the signature of the initialize() method, if it is written in Python - argspec = inspect.getfullargspec(micro_problem_cls.initialize) - if ( - len(argspec.args) == 1 - ): # The first argument in the signature is self - sim_requires_init_data = False - elif len(argspec.args) == 2: - sim_requires_init_data = True - else: - raise Exception( - "The initialize() method of the Micro simulation has an incorrect number of arguments." + first_id = active_sim_lids[0] # First active simulation ID + micro_sims_to_init = ( + active_sim_lids # Only active simulations will be initialized ) - except TypeError: - self._logger.log_info_rank_zero( - "The signature of initialize() method of the micro simulation cannot be determined. Trying to determine the signature by calling the method." - ) - # Try to get the signature of the initialize() method, if it is not written in Python - try: # Try to call the initialize() method without initial data - self._micro_sims[first_id].initialize() - sim_requires_init_data = False - except TypeError: - self._logger.log_info_rank_zero( - "The initialize() method of the micro simulation has arguments. Attempting to call it again with initial data." - ) - try: # Try to call the initialize() method with initial data - self._micro_sims[first_id].initialize(initial_data[first_id]) - sim_requires_init_data = True - except TypeError: - raise Exception( - "The initialize() method of the Micro simulation has an incorrect number of arguments." - ) + are_there_sims_to_init = True + + test_instance = self._model_manager.get_instance( + self._global_number_of_sims + 1, micro_problem_cls + ) + test_data = None + if is_initial_data_available: + test_data = initial_macro_data[0] + ( + self._micro_sims_init, + sim_requires_init_data, + ) = micro_problem_cls.check_initialize( + test_instance, + test_data, + ) + test_instance.destroy() + del test_instance if sim_requires_init_data and not is_initial_data_available: raise Exception( "The initialize() method of the Micro simulation requires initial data, but no initial macro data has been provided." ) - # Get initial data from micro simulations if initialize() method exists - if self._micro_sims_init: + initial_micro_data: dict[str, list] = dict() + if are_there_sims_to_init and self._micro_sims_init: # Call initialize() method of the micro simulation to check if it returns any initial data if sim_requires_init_data: initial_micro_output = self._micro_sims[first_id].initialize( - initial_data[first_id] + initial_macro_data[first_id] ) else: initial_micro_output = self._micro_sims[first_id].initialize() - if ( - initial_micro_output is None - ): # Check if the detected initialize() method returns any data + if initial_micro_output is None: self._logger.log_warning_rank_zero( "The initialize() call of the Micro simulation has not returned any initial data." " This means that the initialize() call has no effect on the adaptivity. The initialize method will nevertheless still be called." @@ -721,14 +736,12 @@ def initialize(self) -> None: if sim_requires_init_data: for i in micro_sims_to_init: - self._micro_sims[i].initialize(initial_data[i]) + self._micro_sims[i].initialize(initial_macro_data[i]) else: for i in micro_sims_to_init: self._micro_sims[i].initialize() else: # Case where the initialize() method returns data if self._is_adaptivity_on: - initial_micro_data: dict[str, list] = dict() - for name in initial_micro_output.keys(): initial_micro_data[name] = [0] * self._local_number_of_sims # Save initial data from first micro simulation as we anyway have it @@ -749,7 +762,7 @@ def initialize(self) -> None: if sim_requires_init_data: for i in micro_sims_to_init: initial_micro_output = self._micro_sims[i].initialize( - initial_data[i] + initial_macro_data[i] ) for name in self._adaptivity_micro_data_names: self._data_for_adaptivity[name][ @@ -764,36 +777,59 @@ def initialize(self) -> None: i ] = initial_micro_output[name] initial_micro_data[name][i] = initial_micro_output[name] - - if ( - self._lazy_init - ): # If lazy initialization is on, initial states of inactive simulations need to be determined - self._adaptivity_controller.get_full_field_micro_output( - initial_micro_data - ) - for i in range(self._local_number_of_sims): - for name in self._adaptivity_micro_data_names: - self._data_for_adaptivity[name][i] = initial_micro_data[ - name - ][i] - del initial_micro_data # Once the initial data is fed into the adaptivity data, it is no longer required - - else: + else: # If adaptivity is off, the returned initial data from the initialize() method will be ignored self._logger.log_warning_rank_zero( "The initialize() method of the Micro simulation returns initial data, but adaptivity is turned off. The returned data will be ignored. The initialize method will nevertheless still be called." ) if sim_requires_init_data: for i in range(1, self._local_number_of_sims): - self._micro_sims[i].initialize(initial_data[i]) + self._micro_sims[i].initialize(initial_macro_data[i]) else: for i in range(1, self._local_number_of_sims): self._micro_sims[i].initialize() - self._micro_sims_have_output = False - if hasattr(micro_problem_cls, "output") and callable( - getattr(micro_problem_cls, "output") - ): - self._micro_sims_have_output = True + self._micro_sims_have_output = micro_problem_cls.check_output() + + self.load_balancing = create_load_balancer( + self._participant, + self._model_manager, + self._adaptivity_controller if self._is_adaptivity_on else None, + self.state_loader, + self.state_setter, + self._logger, + self._config, + self._micro_sims, + self._global_ids_of_local_sims, + self._global_number_of_sims, + self._comm, + self._rank, + ) + + # If lazy initialization is on, initial states of inactive simulations need to be determined + if self._lazy_init: + # If there is initial micro data, and if this rank has sims to init, then the data is to be gathered + if initial_micro_data and are_there_sims_to_init: + initial_micro_data_list: list[dict] = [ + dict(zip(initial_micro_data, t)) + for t in zip(*initial_micro_data.values()) + ] + else: + # Ranks without active simulations provide empty dicts + initial_micro_data_list: list[dict] = [ + dict() for _ in range(self._local_number_of_sims) + ] + + initial_micro_data_list = ( + self._adaptivity_controller.get_full_field_micro_output( + initial_micro_data_list + ) + ) + + for i in range(self._local_number_of_sims): + for name in self._adaptivity_micro_data_names: + self._data_for_adaptivity[name][i] = initial_micro_data_list[i][ + name + ] self._participant.stop_last_profiling_section() @@ -817,7 +853,7 @@ def _read_data_from_precice(self, dt) -> list: """ read_data: dict[str, list] = dict() - if self._is_adaptivity_with_load_balancing: + if self._is_load_balancing: read_vertex_ids = self._global_ids_of_local_sims else: read_vertex_ids = self._mesh_vertex_ids @@ -849,7 +885,7 @@ def _write_data_to_precice(self, data: list) -> None: data : list List of dicts in which keys are names of data and the values are the data to be written to preCICE. """ - if self._is_adaptivity_with_load_balancing: + if self._is_load_balancing: write_vertex_ids = self._global_ids_of_local_sims else: write_vertex_ids = self._mesh_vertex_ids @@ -876,7 +912,9 @@ def _write_data_to_precice(self, data: list) -> None: self._macro_mesh_name, dname, [], np.array([]) ) - def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: + def _solve_micro_simulations( + self, micro_sims_input: list, dt: float, computed_outputs: dict = {} + ) -> list: """ Solve all micro simulations and assemble the micro simulations outputs in a list of dicts format. @@ -887,6 +925,8 @@ def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: solve a micro simulation. dt : float Time step size. + computed_outputs : dict + Dictionary of global ids to already computed outputs Returns ------- @@ -896,11 +936,23 @@ def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: micro_sims_output: list[dict] = [None] * self._local_number_of_sims for count, sim in enumerate(self._micro_sims): + # skip already computed outputs + gid = self._global_ids_of_local_sims[count] + if gid in computed_outputs: + micro_sims_output[count] = computed_outputs[gid] + continue + # If micro simulation has not crashed in a previous iteration, attempt to solve it if not self._has_sim_crashed[count]: # Attempt to solve the micro simulation try: + self.load_balancing.pre_sim_solve( + self._global_ids_of_local_sims[count] + ) micro_sims_output[count] = sim.solve(micro_sims_input[count], dt) + self.load_balancing.post_sim_solve( + self._global_ids_of_local_sims[count] + ) # If simulation crashes, log the error and keep the output constant at the previous iteration's output except Exception as error_message: @@ -945,7 +997,7 @@ def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: return micro_sims_output def _solve_micro_simulations_with_adaptivity( - self, micro_sims_input: list, dt: float + self, micro_sims_input: list, dt: float, computed_outputs: dict = {} ) -> list: """ Adaptively solve micro simulations and assemble the micro simulations outputs in a list of dicts format. @@ -957,6 +1009,8 @@ def _solve_micro_simulations_with_adaptivity( solve a micro simulation. dt : float Time step size. + computed_outputs : dict + Dictionary of global ids to already computed outputs Returns ------- @@ -969,12 +1023,24 @@ def _solve_micro_simulations_with_adaptivity( # Solve all active micro simulations for lid in active_sim_lids: + # skip already computed outputs + gid = self._global_ids_of_local_sims[lid] + if gid in computed_outputs: + micro_sims_output[lid] = computed_outputs[gid] + continue + # If micro simulation has not crashed in a previous iteration, attempt to solve it if not self._has_sim_crashed[lid]: try: + self.load_balancing.pre_sim_solve( + self._global_ids_of_local_sims[lid] + ) micro_sims_output[lid] = self._micro_sims[lid].solve( micro_sims_input[lid], dt ) + self.load_balancing.post_sim_solve( + self._global_ids_of_local_sims[lid] + ) # Mark the micro sim as active for export micro_sims_output[lid]["Active-State"] = 1 @@ -1057,17 +1123,24 @@ def _solve_micro_simulations_with_model_adaptivity( output = None while self._model_adaptivity_controller.should_iterate(): - self._model_adaptivity_controller.switch_models( - self._mesh_vertex_coords, + switched_lids = self._model_adaptivity_controller.switch_models( + self._mesh_vertex_coords[self._global_ids_of_local_sims], self._t, micro_sims_input, output, self._micro_sims, active_sim_ids, ) - output = solve_variant(micro_sims_input, dt) + computed_outputs = {} + if output is not None: + for lid, out in enumerate(output): + if lid in switched_lids: + continue + gid = self._global_ids_of_local_sims[lid] + computed_outputs[gid] = out + output = solve_variant(micro_sims_input, dt, computed_outputs) self._model_adaptivity_controller.check_convergence( - self._mesh_vertex_coords, + self._mesh_vertex_coords[self._global_ids_of_local_sims], self._t, micro_sims_input, output, diff --git a/micro_manager/micro_simulation.py b/micro_manager/micro_simulation.py index 829c6eb5..c784e39d 100644 --- a/micro_manager/micro_simulation.py +++ b/micro_manager/micro_simulation.py @@ -4,8 +4,645 @@ class MicroSimulation. A global ID member variable is defined for the class Simu created object is uniquely identifiable in a global setting. """ +from abc import ABC, abstractmethod +import inspect +import importlib as ipl -def create_simulation_class(micro_simulation_class, sim_class_name=None): +from .tasking.task import ( + ConstructTask, + ConstructLateTask, + DeleteTask, + InitializeTask, + OutputTask, + SolveTask, + SetStateTask, + GetStateTask, +) + + +class MicroSimulationInterface(ABC): + """ + Abstract base class for micro simulations. Users should inherit from this class + when creating their micro simulation and implement all abstract methods. + + The methods ``initialize`` and ``output`` are optional — override them only if + your simulation needs them. The Micro Manager checks ``requires_initialize()`` + and ``requires_output()`` to decide whether to call them. + + Example usage:: + + from micro_manager import MicroSimulationInterface + + class MicroSimulation(MicroSimulationInterface): + def __init__(self, sim_id: int) -> None: + self._sim_id = sim_id + + def initialize(self, initial_data: dict | None = None) -> dict | None: + pass + + def solve(self, macro_data: dict, dt: float) -> dict: + return {} + + def get_state(self) -> object: + return None + + def set_state(self, state: object) -> None: + pass + + def get_global_id(self) -> int: + return self._sim_id + + def set_global_id(self, global_id: int) -> None: + self._sim_id = global_id + + def output(self) -> None: + pass + """ + + def initialize(self, *args, **kwargs) -> dict | None: + """ + Initialize the micro simulation. Called once before the coupling loop starts. + This method is optional. Override it if your simulation requires initialization. + + Parameters + ---------- + initial_data : dict, optional + Initial data passed from the Micro Manager. + + Returns + ------- + dict or None + Optional initial output data to be used in the adaptivity calculation. + """ + pass + + @abstractmethod + def solve(self, micro_sim_input: dict, dt: float) -> dict: + """ + Solve the micro simulation for one time step. + + Parameters + ---------- + micro_sim_input : dict + Input data from the macro simulation. + dt : float + Time step size. + + Returns + ------- + micro_sim_output : dict + Output data to be passed to the macro simulation. + """ + pass + + @abstractmethod + def get_state(self) -> object: + """ + Return the current state of the micro simulation for checkpointing. + + Returns + ------- + state : object + The current state of the micro simulation. + """ + pass + + @abstractmethod + def set_state(self, state: object) -> None: + """ + Set the state of the micro simulation from a checkpoint. + + Parameters + ---------- + state : object + The state to restore. + """ + pass + + @abstractmethod + def get_global_id(self) -> int: + """ + Return the global ID of this micro simulation instance. + + Returns + ------- + global_id : int + Global ID of the micro simulation. + """ + pass + + @abstractmethod + def set_global_id(self, global_id: int) -> None: + """ + Set the global ID of this micro simulation instance. + + Parameters + ---------- + global_id : int + Global ID to assign. + """ + pass + + def output(self) -> None: + """ + Optional output method called after each solve step. + Override this method if your simulation needs to write output at each step. + """ + pass + + def requires_initialize(self) -> bool: + """ + Return True if this simulation class overrides the ``initialize`` method. + The Micro Manager calls this to determine whether initialization is needed. + + Returns + ------- + requires_initialize : bool + True if ``initialize`` is overridden, False otherwise. + """ + return type(self).initialize is not MicroSimulationInterface.initialize + + def requires_output(self) -> bool: + """ + Return True if this simulation class overrides the ``output`` method. + The Micro Manager calls this to determine whether output is needed. + + Returns + ------- + requires_output : bool + True if ``output`` is overridden, False otherwise. + """ + return type(self).output is not MicroSimulationInterface.output + + def destroy(self) -> None: + """ + Relinquishes allocated resources. This will result in an object with invalid state. + Do not use after this. Object should be scheduled for deletion. + Calling destroy gives a controlled mechanism to release allocation, + contrary to GC. + """ + pass + + +class MicroSimulationLocal(MicroSimulationInterface): + def __init__(self, gid, late_init, sim_cls): + self._gid = gid + self._instance = sim_cls(-1 if late_init else gid) + + def solve(self, micro_sim_input, dt): + return self._instance.solve(micro_sim_input, dt) + + def get_state(self): + return self._instance.get_state() + + def set_state(self, state): + return self._instance.set_state(state) + + def get_global_id(self): + return self._gid + + def set_global_id(self, global_id): + self._gid = global_id + + def __getattr__(self, name): + return getattr(self._instance, name) + + def initialize(self, *args, **kwargs): + return self._instance.initialize(*args, **kwargs) + + def output(self): + return self._instance.output() + + def requires_initialize(self) -> bool: + return self._instance.requires_initialize() + + def requires_output(self) -> bool: + return self._instance.requires_output() + + def destroy(self): + self._instance = None + + +class MicroSimulationRemote(MicroSimulationInterface): + def __init__(self, gid, late_init, num_ranks, conn, cls_path, sim_cls): + self._cls_path = cls_path + self._gid = gid + self._num_ranks = num_ranks + self._conn = conn + self._sim_cls = sim_cls + + construct_cls = ConstructLateTask if late_init else ConstructTask + for worker_id in range(self._num_ranks): + task = construct_cls.send_args(self._gid, self._cls_path) + self._conn.send(worker_id, task) + + for worker_id in range(self._num_ranks): + self._conn.recv(worker_id) + + def solve(self, micro_sim_input, dt): + for worker_id in range(self._num_ranks): + task = SolveTask.send_args(self._gid, micro_sim_input, dt) + self._conn.send(worker_id, task) + + result = None + for worker_id in range(self._num_ranks): + output = self._conn.recv(worker_id) + if worker_id == 0: + result = output + + return result + + def get_state(self): + for worker_id in range(self._num_ranks): + task = GetStateTask.send_args(self._gid) + self._conn.send(worker_id, task) + + result = {} + for worker_id in range(self._num_ranks): + result[worker_id] = self._conn.recv(worker_id) + + return result + + def set_state(self, state): + for worker_id in range(self._num_ranks): + task = SetStateTask.send_args(self._gid, state[worker_id]) + self._conn.send(worker_id, task) + + result = {} + for worker_id in range(self._num_ranks): + result[worker_id] = self._conn.recv(worker_id) + self._gid = result[0] + + def get_global_id(self): + return self._gid + + def set_global_id(self, global_id): + self._gid = global_id + + def initialize(self, *args, **kwargs): + for worker_id in range(self._num_ranks): + task = InitializeTask.send_args(self._gid, *args, **kwargs) + self._conn.send(worker_id, task) + + result = None + for worker_id in range(self._num_ranks): + output = self._conn.recv(worker_id) + if worker_id == 0: + result = output + + return result + + def output(self): + for worker_id in range(self._num_ranks): + task = OutputTask.send_args(self._gid) + self._conn.send(worker_id, task) + + result = None + for worker_id in range(self._num_ranks): + output = self._conn.recv(worker_id) + if worker_id == 0: + result = output + + return result + + def requires_initialize(self) -> bool: + return self._sim_cls.initialize is not MicroSimulationInterface.initialize + + def requires_output(self) -> bool: + return self._sim_cls.output is not MicroSimulationInterface.output + + def destroy(self): + for worker_id in range(self._num_ranks): + task = DeleteTask.send_args(self._gid) + self._conn.send(worker_id, task) + + for worker_id in range(self._num_ranks): + self._conn.recv(worker_id) + + +class MicroSimulationWrapper(MicroSimulationInterface): + """ + If only a single rank is in use: will contain the micro sim instance. + Otherwise, it will delegate method calls to workers and not contain state. + """ + + def __init__(self, name, sim_cls, cls_path, global_id, late_init, num_ranks, conn): + self._impl = None + + if num_ranks > 1 and conn is not None: + self._impl = MicroSimulationRemote( + global_id, late_init, num_ranks, conn, cls_path, sim_cls + ) + else: + self._impl = MicroSimulationLocal(global_id, late_init, sim_cls) + + self._external_data = dict() + self._name = name + + def solve(self, micro_sim_input, dt): + return self._impl.solve(micro_sim_input, dt) + + def get_state(self): + return self._impl.get_state() + + def set_state(self, state): + return self._impl.set_state(state) + + def get_global_id(self): + return self._impl.get_global_id() + + def set_global_id(self, global_id): + return self._impl.set_global_id(global_id) + + def initialize(self, *args, **kwargs): + return self._impl.initialize(*args, **kwargs) + + def output(self): + return self._impl.output() + + def requires_initialize(self) -> bool: + return self._impl.requires_initialize() + + def requires_output(self) -> bool: + return self._impl.requires_output() + + def destroy(self): + return self._impl.destroy() + + def __getattr__(self, name): + return getattr(self._impl, name) + + @property + def attachments(self): + return self._external_data + + @attachments.setter + def attachments(self, value): + self._external_data = value + + @property + def name(self): + return self._name + + +class MicroSimulationClass: + def __init__(self, sim_cls, cls_path, name, num_ranks, conn, log): + self._sim_cls = sim_cls + self._cls_path = cls_path + self._name = name + self._num_ranks = num_ranks + self._conn = conn + self._log = log + + @property + def name(self): + return self._name + + def __call__(self, gid, *, late_init=False): + return MicroSimulationWrapper( + self._name, + self._sim_cls, + self._cls_path, + gid, + late_init, + self._num_ranks, + self._conn, + ) + + @property + def backend_cls(self): + return self._sim_cls + + def check_initialize( + self, test_instance: MicroSimulationInterface, test_input: dict + ) -> tuple[bool, bool]: + """ + Check whether the micro simulation class implements ``initialize``. + + Since ``load_backend_class`` guarantees that ``self._sim_cls`` always + inherits from ``MicroSimulationInterface``, we can rely on + ``requires_initialize()`` directly. No ``issubclass`` guard is needed. + + Parameters + ---------- + test_instance : object + An instance of the micro simulation class used for signature probing. + test_input : dict + Sample input data used to probe whether ``initialize`` accepts arguments. + + Returns + ------- + check_result : tuple[bool, bool] + (has_initialize, requires_initial_data) + """ + if not test_instance.requires_initialize(): + return False, False + + has_args = False + + # Try to get the signature of the initialize() method, if it is written in Python + try: + argspec = inspect.getfullargspec(self._sim_cls.initialize) + # The first argument in the signature is self + if len(argspec.args) == 1: + has_args = False + elif len(argspec.args) == 2: + has_args = True + else: + raise Exception( + "The initialize() method of the Micro simulation has an incorrect number of arguments." + ) + except TypeError: + self._log.log_info_rank_zero( + "The signature of initialize() method of the micro simulation cannot be determined. " + + "Trying to determine the signature by calling the method." + ) + # Try to call the initialize() method without initial data + try: + test_instance.initialize() + has_args = False + except TypeError: + self._log.log_info_rank_zero( + "The initialize() method of the micro simulation has arguments. " + + "Attempting to call it again with initial data." + ) + try: + test_instance.initialize(test_input) + has_args = True + except TypeError: + raise Exception( + "The initialize() method of the Micro simulation has an incorrect number of arguments." + ) + + return True, has_args + + def check_output(self) -> bool: + """ + Check whether the micro simulation class implements ``output``. + + Since ``load_backend_class`` guarantees that ``self._sim_cls`` always + inherits from ``MicroSimulationInterface``, we can rely on + ``requires_output()`` directly at the class level. + + Returns + ------- + check_result : bool + True if the micro simulation class overrides the ``output`` method. + """ + return self._sim_cls.output is not MicroSimulationInterface.output + + +def _wrap_non_interface_class(cls: type, path_to_micro_file: str) -> type: + """ + Dynamically create a class that inherits from MicroSimulationInterface + and delegates all method calls to the provided class. + + This ensures that load_backend_class always returns a class that adheres + to MicroSimulationInterface, even for pybind11 classes or legacy classes + that do not explicitly inherit from it. + + Parameters + ---------- + cls : type + The original micro simulation class (e.g. loaded via pybind11). + path_to_micro_file : str + Path string used for the deprecation warning message. + + Returns + ------- + type + A new class inheriting from MicroSimulationInterface that wraps cls. + """ + import warnings + + warnings.warn( + "The MicroSimulation class in '{}' does not inherit from MicroSimulationInterface. " + "Please update your class definition to: " + "class MicroSimulation(MicroSimulationInterface). " + "In a future version this will become an error.".format(path_to_micro_file), + DeprecationWarning, + stacklevel=3, + ) + + # Determine whether the original class provides initialize / output + has_initialize = callable(getattr(cls, "initialize", None)) + has_output = callable(getattr(cls, "output", None)) + + # Build the class body: __init__ and mandatory interface methods + class_body = """ +def __init__(self, global_id): + self._wrapped = wrapped_cls(global_id) + +def solve(self, micro_sim_input, dt): + return self._wrapped.solve(micro_sim_input, dt) + +def get_state(self): + return self._wrapped.get_state() + +def set_state(self, state): + return self._wrapped.set_state(state) + +def get_global_id(self): + return self._wrapped.get_global_id() + +def set_global_id(self, global_id): + self._wrapped.set_global_id(global_id) + +def __getattr__(self, name): + return getattr(self._wrapped, name) +""" + + # Only add initialize override if the wrapped class actually has it, + # so that requires_initialize() returns True for those classes. + if has_initialize: + class_body += """ +def initialize(self, *args, **kwargs): + return self._wrapped.initialize(*args, **kwargs) +""" + + # Only add output override if the wrapped class actually has it, + # so that requires_output() returns True for those classes. + if has_output: + class_body += """ +def output(self): + return self._wrapped.output() +""" + + class_dict = {} + exec(class_body, {"wrapped_cls": cls, "__builtins__": __builtins__}, class_dict) + + wrapper_cls = type( + "CompatibilityWrapper_{}".format(cls.__name__), + (MicroSimulationInterface,), + class_dict, + ) + return wrapper_cls + + +def load_backend_class(path_to_micro_file: str) -> type: + """ + Load the MicroSimulation class from the given module path. + + Always returns a class that inherits from MicroSimulationInterface. + If the loaded class does not inherit from it (e.g. pybind11 classes or + legacy classes), it is wrapped in a dynamically created adapter class + that delegates all calls to the original and correctly implements + requires_initialize() and requires_output(). + + Parameters + ---------- + path_to_micro_file : str + Dotted module path to the micro simulation file. + + Returns + ------- + type + A class inheriting from MicroSimulationInterface. + """ + + def try_load(name): + try: + return getattr(ipl.import_module(path_to_micro_file, name), name) + except ImportError as ie: + return None + except AttributeError as ae: + return None + + def check_cls(cls): + try: + inherits = issubclass(cls, MicroSimulationInterface) + except TypeError: + # pybind11 extension types may not support issubclass — wrap them + inherits = False + + if not inherits: + cls = _wrap_non_interface_class(cls, path_to_micro_file) + return cls + + CLS_NAME = "MicroSimulation" + # attempt to load with base name + result = try_load(CLS_NAME) + if result is not None: + return check_cls(result) + + # attempt to load with appended indices + for i in range(10): + result = try_load(f"{CLS_NAME}{i}") + if result is not None: + return check_cls(result) + + # failed to load any class + raise RuntimeError(f"Could not load micro simulation from {path_to_micro_file}") + + +def create_simulation_class( + log, + micro_simulation_class, + path_to_micro_file, + num_ranks, + conn=None, + sim_class_name=None, +): """ Creates a class Simulation which inherits from the class of the micro simulation. @@ -22,6 +659,15 @@ def create_simulation_class(micro_simulation_class, sim_class_name=None): Simulation : class Definition of class Simulation defined in this function. """ + if not hasattr(micro_simulation_class, "get_global_id"): + raise ValueError("Invalid micro simulation class") + if not hasattr(micro_simulation_class, "get_state"): + raise ValueError("Invalid micro simulation class") + if not hasattr(micro_simulation_class, "set_state"): + raise ValueError("Invalid micro simulation class") + if not hasattr(micro_simulation_class, "solve"): + raise ValueError("Invalid micro simulation class") + if sim_class_name is None: if not hasattr(create_simulation_class, "sim_id"): create_simulation_class.sim_id = 0 @@ -29,21 +675,7 @@ def create_simulation_class(micro_simulation_class, sim_class_name=None): create_simulation_class.sim_id += 1 sim_class_name = f"MicroSimulation{create_simulation_class.sim_id}" - sim_class_body = """ -def __init__(self, global_id): - micro_simulation_class.__init__(self, global_id) - self._global_id = global_id - -def get_global_id(self) -> int: - return self._global_id - """ - sim_class_dict = {} - local_globals = { - "__builtins__": __builtins__, - "micro_simulation_class": micro_simulation_class, - } - exec(sim_class_body, local_globals, sim_class_dict) - # print(sim_class_dict) - sim_class = type(sim_class_name, (micro_simulation_class,), sim_class_dict) - - return sim_class + result_cls = MicroSimulationClass( + micro_simulation_class, path_to_micro_file, sim_class_name, num_ranks, conn, log + ) + return result_cls diff --git a/micro_manager/model_manager.py b/micro_manager/model_manager.py new file mode 100644 index 00000000..34927aeb --- /dev/null +++ b/micro_manager/model_manager.py @@ -0,0 +1,223 @@ +from micro_manager.micro_simulation import ( + MicroSimulationClass, + MicroSimulationWrapper, + MicroSimulationInterface, +) + + +class ModelWrapper(MicroSimulationInterface): + """ + Stateless Model Wrapper, will delegate any method call to the main compute instance. + This is used to replace instances in the main simulation container. + """ + + def __init__(self, global_id, backend): + self._global_id = global_id + self._backend = backend + + def set_global_id(self, global_id): + self._global_id = global_id + + def get_global_id(self) -> int: + return self._global_id + + def solve(self, macro_data, dt): + return self._backend.solve(macro_data, dt) + + def get_state(self): + return self._backend.get_state() + + def set_state(self, state): + self._backend.set_state(state) + + def initialize(self, *args, **kwargs): + return self._backend.initialize(*args, **kwargs) + + def output(self): + return self._backend.output() + + def destroy(self): + # we do not want to delete our compute instance + pass + + @property + def __class__(self): + return self._backend.__class__ + + @property + def attachments(self): + return self._backend.attachments + + @attachments.setter + def attachments(self, value): + self._backend.attachments = value + + @property + def name(self): + return self._backend.name + + +class ModelManager: + """ + Manages all used micro simulation models. Stores their classes and checks whether they may + use model instancing. To generate instances use the get_instance method regardless of model instancing, + as the ModelManager handles either case. + """ + + def __init__(self): + self._registered_classes: list[MicroSimulationClass] = [] + self._stateless_map: dict[MicroSimulationClass, bool] = dict() + self._backend_map: dict[MicroSimulationClass, MicroSimulationWrapper] = dict() + + def register(self, micro_sim_cls: MicroSimulationClass, stateless: bool): + """ + Register a micro simulation class to create an instance of it later. + + Parameters + ---------- + micro_sim_cls : MicroSimulationClass + Micro simulation class to register. + stateless: bool + Is the simulation class stateless. + """ + if micro_sim_cls in self._registered_classes: + return + + self._registered_classes.append(micro_sim_cls) + self._stateless_map[micro_sim_cls] = stateless + + if stateless: + self._backend_map[micro_sim_cls] = micro_sim_cls( + len(self._registered_classes) - 1 + ) + + def get_cls_by_name(self, name: str) -> MicroSimulationClass: + """ + Returns the class determined by its name + + Parameters + ---------- + name: str + name of registered class + + Returns + ------- + sim_class: MicroSimulationClass + Class given by name + """ + cls_names = [cls.name for cls in self._registered_classes] + idx = cls_names.index(name) + return self._registered_classes[idx] + + def get_cls_by_idx(self, idx: int) -> MicroSimulationClass: + """ + Returns the class determined by its index + + Parameters + ---------- + idx: int + index of registered class + + Returns + ------- + sim_class: MicroSimulationClass + Class given by index + """ + return self._registered_classes[idx] + + def is_stateless(self, name: str) -> bool: + """ + Returns whether the class given by its name is stateless. + + Parameters + ---------- + name: str + name of registered class + + Returns + ------- + is_stateless: bool + true if class is stateless + """ + cls = self.get_cls_by_name(name) + return self._stateless_map[cls] + + def get_instance_by_name( + self, gid: int, name: str, *, late_init: bool = False + ) -> MicroSimulationInterface: + """ + Creates an instance of the requested class determined by its name. If the class should be initialized later, + the request will be delegated to the micro simulation object (in case it supports it). + + Parameters + ---------- + gid: int + Global Simulation ID + name: str + Requested micro simulation class + late_init: bool + Should the simulation be initialized later? + + Returns + ------- + micro_sim : MicroSimulationInterface + Instance of the requested micro simulation class, either delegator or compute instance + """ + return self.get_instance(gid, self.get_cls_by_name(name), late_init=late_init) + + def get_instance_by_idx( + self, gid: int, idx: int, *, late_init: bool = False + ) -> MicroSimulationInterface: + """ + Creates an instance of the requested class determined by its index. If the class should be initialized later, + the request will be delegated to the micro simulation object (in case it supports it). + + Parameters + ---------- + gid: int + Global Simulation ID + idx: int + Index of requested micro simulation class + late_init: bool + Should the simulation be initialized later? + + Returns + ------- + micro_sim : MicroSimulationInterface + Instance of the requested micro simulation class, either delegator or compute instance + """ + return self.get_instance( + gid, self._registered_classes[idx], late_init=late_init + ) + + def get_instance( + self, gid: int, micro_sim_cls: MicroSimulationClass, *, late_init: bool = False + ) -> MicroSimulationInterface: + """ + Creates an instance of the requested class. If the class should be initialized later, + the request will be delegated to the micro simulation object (in case it supports it). + + Parameters + ---------- + gid: int + Global Simulation ID + micro_sim_cls: MicroSimulationClass + Requested micro simulation class + late_init: bool + Should the simulation be initialized later? + + Returns + ------- + micro_sim : MicroSimulationInterface + Instance of the requested micro simulation class, either delegator or compute instance + """ + if micro_sim_cls not in self._registered_classes: + raise RuntimeError("Trying to create instance of unknown class!") + + if self._stateless_map[micro_sim_cls]: + return ModelWrapper( + gid, + self._backend_map[micro_sim_cls], + ) + else: + return micro_sim_cls(gid, late_init=late_init) diff --git a/micro_manager/snapshot/snapshot.py b/micro_manager/snapshot/snapshot.py index 4b6ae71f..618e5987 100644 --- a/micro_manager/snapshot/snapshot.py +++ b/micro_manager/snapshot/snapshot.py @@ -17,7 +17,7 @@ from micro_manager.micro_manager import MicroManager from .dataset import ReadWriteHDF -from micro_manager.micro_simulation import create_simulation_class +from micro_manager.micro_simulation import create_simulation_class, load_backend_class from micro_manager.tools.logging_wrapper import Logger @@ -84,7 +84,13 @@ def solve(self) -> None: - Merge output in parallel run. """ - micro_problem_cls = create_simulation_class(self._micro_problem) + micro_problem_cls = create_simulation_class( + self._logger, + self._micro_problem, + self._config.get_micro_file_name(), + 1, + None, + ) # Loop over all macro parameters for elems in range(self._local_number_of_sims): @@ -256,12 +262,7 @@ def initialize(self) -> None: for i in range(self._local_number_of_sims): self._global_ids_of_local_sims.append(sim_id) sim_id += 1 - self._micro_problem = getattr( - importlib.import_module( - self._config.get_micro_file_name(), "MicroSimulation" - ), - "MicroSimulation", - ) + self._micro_problem = load_backend_class(self._config.get_micro_file_name()) self._micro_sims_have_output = False if hasattr(self._micro_problem, "output") and callable( diff --git a/micro_manager/tasking/__init__.py b/micro_manager/tasking/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/micro_manager/tasking/connection.py b/micro_manager/tasking/connection.py new file mode 100644 index 00000000..881a67e8 --- /dev/null +++ b/micro_manager/tasking/connection.py @@ -0,0 +1,444 @@ +import pickle +import psutil +import socket +import struct +import subprocess +import os +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional +from mpi4py import MPI + +# worker_main gets spawned not part of the micro_manager package. +# Therefore, imports do not work properly. As worker_main requires connection.py, +# this here is a workaround to still load it. +import sys +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parent)) +from task import RegisterAllTask, ShutdownTask + + +class Connection(ABC): + @abstractmethod + def send(self, dst_id: int, obj: Any) -> None: + pass + + @abstractmethod + def recv(self, src_id: int) -> Any: + pass + + @abstractmethod + def close(self) -> None: + pass + + @abstractmethod + def is_open(self) -> bool: + pass + + +class MPIConnection(Connection): + def __init__(self): + self.inter_comm = None + self._num_workers = 0 + self.open = False + + @classmethod + def create_workers( + cls, worker_exec: str, mpi_args: Optional, n_workers: int + ) -> "MPIConnection": + args = [worker_exec] + args.extend(mpi_args or []) + + env = os.environ.copy() + check_btl_base = ( + "OMPI_MCA_btl" not in env or env["OMPI_MCA_btl"] != "self,vader,tcp" + ) + check_btl_tcp = ( + "OMPI_MCA_btl_tcp_if_include" not in env + or env["OMPI_MCA_btl_tcp_if_include"] != "lo" + ) + check_mca_oob = ( + "OMPI_MCA_oob_tcp_if_include" not in env + or env["OMPI_MCA_oob_tcp_if_include"] != "lo" + ) + if check_btl_base or check_btl_tcp or check_mca_oob: + msg = ( + "Cannot launch MPI workers. Please set the following environment variables:\n" + "\tOMPI_MCA_btl=self,vader,tcp\n" + "\tOMPI_MCA_btl_tcp_if_include=lo\n" + "\tOMPI_MCA_oob_tcp_if_include=lo\n" + ) + raise RuntimeError(msg) + + comm = MPI.COMM_SELF + conn = cls() + conn._num_workers = n_workers + conn.inter_comm = comm.Spawn( + "python", + args=args, + maxprocs=n_workers, + ) + conn.open = True + return conn + + @classmethod + def connect_to_micro_manager(cls, parent_comm) -> "MPIConnection": + conn = cls() + conn.inter_comm = parent_comm + conn.open = True + return conn + + def send(self, dst_id: int, obj: Any) -> None: + """ + Sends any data (obj) to the worker with id (dst_id). dst_id corresponds to the + worker MPI process rank. Data is encoded by pickling it. Thus, one form of pickling must be implemented for + the data type of (obj). Data is transferred using an MPI inter-process communicator. + + Parameters + ---------- + dst_id : int + Worker MPI process rank. + obj : Any + Data to send. (needs implemented pickling interface) + """ + if not self.open: + return + data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + self.inter_comm.send(data, dest=dst_id, tag=0) + + def recv(self, src_id: int) -> Any: + """ + Receives data from the worker with id (src_id). src_id corresponds to the + worker MPI process rank. Data is transferred using an MPI inter-process communicator. + + Parameters + ---------- + src_id : int + Worker MPI process rank. + """ + if not self.open: + return None + data = self.inter_comm.recv(source=src_id, tag=0) + return pickle.loads(data) + + def close(self) -> None: + if self._num_workers > 0: + for i in range(self._num_workers): + self.send(i, ShutdownTask.send_args()) + self.open = False + self.inter_comm.Disconnect() + + def is_open(self) -> bool: + return self.open + + +class SocketConnection(Connection): + def __init__(self): + self.sockets: Dict[int, socket.socket] = {} + self.open = False + + @classmethod + def create_workers( + cls, worker_exec: str, launcher: list, host: str, n_workers: int, env_opts: dict + ) -> "SocketConnection": + # create listening socket with ephemeral port + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.bind((host, 0)) # kernel picks free port + server.listen() + port = server.getsockname()[1] + + executable = [ + "python", + worker_exec, + "--backend", + "socket", + "--host", + host, + "--port", + str(port), + ] + cmd = [] + cmd.extend(launcher) + cmd.extend(executable) + + env = os.environ.copy() + env.update(env_opts) + subprocess.Popen(cmd, env=env) + + conn = cls() + for wid in range(n_workers): + sock, _ = server.accept() + conn.sockets[wid] = sock + + server.close() + conn.open = True + return conn + + @classmethod + def connect_to_micro_manager( + cls, worker_id: int, host: str, port: int + ) -> "SocketConnection": + conn = cls() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + conn.sockets[worker_id] = sock + conn.open = True + return conn + + def send(self, dst_id: int, obj: Any) -> None: + """ + Sends any data (obj) to the worker with id (dst_id). dst_id corresponds to the + worker MPI process rank. Data is encoded by pickling it. Thus, one form of pickling must be implemented for + the data type of (obj). Data is transferred using a socket by first encoding a header + (containing the data size), followed by the actual data. + + Parameters + ---------- + dst_id : int + Worker MPI process rank. + obj : Any + Data to send. (needs implemented pickling interface) + """ + if not self.open: + return + sock = self.sockets[dst_id] + data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + header = struct.pack("!Q", len(data)) + sock.sendall(header + data) + + def recv(self, src_id: int) -> Any: + """ + Receives data from the worker with id (src_id). src_id corresponds to the + worker MPI process rank. Data is transferred using a socket. First reads the header to determine the data size. + Then reads the incoming data. + + Parameters + ---------- + src_id : int + Worker MPI process rank. + """ + if not self.open: + return None + sock = self.sockets[src_id] + header = sock.recv(8) + if not header: + raise EOFError + (size,) = struct.unpack("!Q", header) + payload = b"" + while len(payload) < size: + chunk = sock.recv(size - len(payload)) + if not chunk: + raise EOFError + payload += chunk + return pickle.loads(payload) + + def close(self) -> None: + for sock in self.sockets.values(): + sock.close() + self.open = False + self.sockets.clear() + + def is_open(self) -> bool: + return self.open + + +def get_mpi_pinning(mpi_impl: str, num_workers: int, hostfile: str): + """ + Returns a list containing args to determine MPI process pinning depending on the MPI implementation. + + Parameters + ---------- + mpi_impl : string + MPI implementation + + Returns + ------- + list + pinning args + """ + args = [] + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + + options = {} + if mpi_impl == "intel": + if os.path.exists(hostfile): + with open(hostfile, "r") as f: + hosts = f.readlines() + else: + raise RuntimeError("Cannot determine target nodes") + + if size % len(hosts) != 0: + raise RuntimeError("Number of ranks must be divisible by number of hosts") + + mm_ppn = size // len(hosts) + node_idx = rank // mm_ppn + node = hosts[node_idx].replace("\n", "") + + locations_int = list(os.sched_getaffinity(0)) + locations = ",".join([str(i) for i in locations_int]) + + # See: https://www.intel.com/content/www/us/en/developer/tools/oneapi/mpi-library-pinning-simulator.html + # for more details and a nice visualization + options.update( + { + "I_MPI_DEBUG": "5", + "I_MPI_PIN": "1", + "I_MPI_PIN_CELL": "core", + "I_MPI_PIN_DOMAIN": "1", + "I_MPI_PIN_PROCESSOR_LIST": locations, + } + ) + + for key, value in options.items(): + args.append("-genv") + args.append(f"{key}={value}") + args.append("-host") + args.append(f"{node}") + + if mpi_impl == "open": + args.extend(["--bind-to", "core"]) + locations = ",".join([str(i + rank * num_workers) for i in range(num_workers)]) + args.extend(["--map-by", f"PE-LIST={locations}:ORDERED"]) + args.extend(["--report-bindings"]) + + return args, options + + +def get_local_ip(preferred_ifaces=None) -> str: + """ + Returns a non-loopback IPv4 address without accessing external networks. + + Parameters + ---------- + preferred_ifaces : list[str], optional + If provided, try interfaces in this order first (e.g., ["ib0", "eno1"]) + + Returns + ------- + str + The selected IPv4 address + """ + addrs = psutil.net_if_addrs() + + candidates = [] + + # Iterate over preferred interfaces first + if preferred_ifaces: + for name in preferred_ifaces: + if name not in addrs: + continue + for a in addrs[name]: + # trying to find an interface that is not loopback (127.xx) + if a.family == socket.AF_INET and not a.address.startswith("127."): + return a.address + + # Fallback: iterate all interfaces + for name, iface_addrs in addrs.items(): + for a in iface_addrs: + if a.family == socket.AF_INET: + ip = a.address + # trying to find an interface that is not loopback (127.xx) or link-local (169.254.xx) + if not ip.startswith("127.") and not ip.startswith("169.254."): + candidates.append(ip) + + if candidates: + return candidates[0] + + raise RuntimeError("No non-loopback IPv4 address found") + + +def spawn_local_workers( + worker_exec: str, + n_workers: int, + backend: str, + is_slurm: bool, + mpi_impl: str, + hostfile: str, +): + """ + Spawn worker processes. On Slurm systems: MPI spawn now supported, socket backend enforced. + Ephemeral port auto-selected. + + Parameters + ---------- + worker_exec : str + path to worker executable + n_workers : int + number of worker processes, must be > 1 otherwise returns None + backend : str + mpi or socket + is_slurm : bool + is our system slurm based? + mpi_impl : string + MPI implementation [intel or open] + hostfile : str + Path to Hostfile containing hosts for workers + + Returns + ------- + conn : Connection + Established connection on generator side + """ + if n_workers <= 1: + return None + conn = None + + # MPI BACKEND (non-Slurm only) + if backend == "mpi": + if is_slurm: + raise RuntimeError( + "MPI backend is not supported under Slurm. " + "Use socket backend instead." + ) + comm = MPI.COMM_WORLD + local_rank = comm.Get_rank() + conn = MPIConnection.create_workers( + worker_exec=worker_exec, + mpi_args=[ + "--backend", + "mpi", + ], + n_workers=n_workers, + ) + + # SOCKET BACKEND + if backend == "socket": + host = get_local_ip() + pin_args, pin_options = get_mpi_pinning(mpi_impl, n_workers, hostfile) + # launch workers + launcher = None + if is_slurm: + launcher = [ + "srun", + f"--ntasks={n_workers}", + "--cpus-per-task=1", + "--cpu-bind=cores", + "--exclusive", + ] + else: + launcher = ["mpiexec"] + if mpi_impl == "intel": + launcher.extend(["-ppn", str(n_workers)]) + launcher.extend(["-n", str(n_workers)]) + launcher.extend(pin_args) + + conn = SocketConnection.create_workers( + worker_exec=worker_exec, + launcher=launcher, + host=host, + n_workers=n_workers, + env_opts=pin_options, + ) + + from ..micro_simulation import load_backend_class + + # Send RegisterAllTask to all workers, which sets up the local worker state + # and registers all potentially used tasks on the workers side. + # By doing so, less pickling needs to be done during operation. Only the task name and data need to be transferred. + # Workers can then locally re-construct the task based on the registered tasks. + for worker_id in range(n_workers): + conn.send(worker_id, RegisterAllTask(load_backend_class)) + conn.recv(worker_id) + + return conn diff --git a/micro_manager/tasking/task.py b/micro_manager/tasking/task.py new file mode 100644 index 00000000..9e666b1d --- /dev/null +++ b/micro_manager/tasking/task.py @@ -0,0 +1,228 @@ +class Task: + """ + This is the general task interface. + Each task is callable and will be provided with a global state object (state_data). + + Inheriting classes may define functions to be executed that require more args than just the state object. + These will be passed on as args and kwargs. Args and kwargs are bound to the task object upon construction. + Thus, each call to a task object only requires the state object. + + Inheriting classes need to call super.__init__ with the function to be called and the args and kwargs to be bound. + """ + + def __init__(self, fn, *args, **kwargs): + self.fn = fn + self.args = args + self.kwargs = kwargs + + def __call__(self, state_data: dict): + return self.fn(*self.args, state_data=state_data, **self.kwargs) + + @classmethod + def send_args(cls, *args, **kwargs): + """ + Used to get a representation of the task without the need to pickle the task. + """ + return cls.__name__, args, kwargs + + +class DeleteTask(Task): + def __init__(self, gid): + super().__init__(DeleteTask.delete_gid, gid=gid) + + @staticmethod + def delete_gid(gid, state_data): + if gid in state_data["sim_instances"]: + del state_data["sim_instances"][gid] + return None + + +class ConstructTask(Task): + """ + Construction Task: Given a gid and micro simulation class path, it will construct an instance and store it + in the state object under ['sim_instances'][gid]. If the desired class has not yet been loaded, then this + will be done by the 'load_function' prior to construction. + """ + + def __init__(self, gid, cls_path): + super().__init__(ConstructTask.initializer, gid=gid, cls_path=cls_path) + + @staticmethod + def initializer(gid, cls_path, state_data): + if cls_path not in state_data["sim_classes"]: + import os + import sys + from pathlib import Path + + ms_dir = os.path.abspath(str(Path(cls_path).resolve().parent)) + sys.path.append(ms_dir) + _, file_name = os.path.split(os.path.abspath(str(Path(cls_path).resolve()))) + + state_data["sim_classes"][cls_path] = state_data["load_function"](file_name) + cls = state_data["sim_classes"][cls_path] + + if gid in state_data["sim_instances"]: + del state_data["sim_instances"][gid] + state_data["sim_instances"][gid] = cls(gid) + return None + + +class ConstructLateTask(Task): + """ + Similar to ConstructTask, it will construct an instance and store it. However, it will pass -1 as the gid to the + instance to allow for late initialization, if the micro simulation supports it. + """ + + def __init__(self, gid, cls_path): + super().__init__(ConstructLateTask.initializer, gid=gid, cls_path=cls_path) + + @staticmethod + def initializer(gid, cls_path, state_data): + if cls_path not in state_data["sim_classes"]: + import os + import sys + from pathlib import Path + + ms_dir = os.path.abspath(str(Path(cls_path).resolve().parent)) + sys.path.append(ms_dir) + _, file_name = os.path.split(os.path.abspath(str(Path(cls_path).resolve()))) + + state_data["sim_classes"][cls_path] = state_data["load_function"](file_name) + cls = state_data["sim_classes"][cls_path] + + if gid in state_data["sim_instances"]: + del state_data["sim_instances"][gid] + state_data["sim_instances"][gid] = cls(-1) + return None + + +class SolveTask(Task): + """ + Given a gid, input data and the current dt, the SolveTask will call the solve function of its respective + simulation object, that is stored in the state object under ['sim_instances'][gid]. + """ + + def __init__(self, gid, sim_input, dt): + super().__init__(SolveTask.solve, gid=gid, sim_input=sim_input, dt=dt) + + @staticmethod + def solve(gid, sim_input, dt, state_data): + sim_output = state_data["sim_instances"][gid].solve(sim_input, dt) + return sim_output + + +class GetStateTask(Task): + """ + Given a gid, the GetStateTask will call the get_state function of its respective + simulation object, that is stored in the state object under ['sim_instances'][gid]. + """ + + def __init__(self, gid): + super().__init__(GetStateTask.get, gid=gid) + + @staticmethod + def get(gid, state_data): + return state_data["sim_instances"][gid].get_state() + + +class SetStateTask(Task): + """ + Given a gid and a state the SetStateTask will call the set_state function of its respective + simulation object, that is stored in the state object under ['sim_instances'][gid]. + """ + + def __init__(self, gid, state): + super().__init__(SetStateTask.set, gid=gid, state=state) + + @staticmethod + def set(gid, state, state_data): + state_data["sim_instances"][gid].set_state(state) + + # if gid was changed, we want to move it to the right location + check_gid = state_data["sim_instances"][gid].get_global_id() + if check_gid != gid: + state_data["sim_instances"][check_gid] = state_data["sim_instances"][gid] + del state_data["sim_instances"][gid] + return check_gid + + return gid + + +class InitializeTask(Task): + """ + Given a gid and arbitrary arguments the InitializeTask will call the initialize function of its respective + simulation object, that is stored in the state object under ['sim_instances'][gid]. + All arguments will be passed along. + """ + + def __init__(self, gid, *args, **kwargs): + super().__init__(InitializeTask.initialize, *args, gid=gid, **kwargs) + + @staticmethod + def initialize(gid, state_data, *args, **kwargs): + return state_data["sim_instances"][gid].initialize(*args, **kwargs) + + +class OutputTask(Task): + """ + Given a gid, the OutputTask will call the output function of its respective + simulation object, that is stored in the state object under ['sim_instances'][gid]. + """ + + def __init__(self, gid): + super().__init__(OutputTask.output, gid=gid) + + @staticmethod + def output(gid, state_data): + return state_data["sim_instances"][gid].output() + + +class ShutdownTask(Task): + """ + The ShutdownTask will raise an exception in order to exit out of the work loop with in the worker_main. + """ + + def __init__(self): + super().__init__(ShutdownTask.shutdown) + + @staticmethod + def shutdown(state_data): + raise RuntimeError("Stopping Worker") + + +class RegisterAllTask(Task): + """ + Sets up the local worker state and registers all potentially used tasks on the workers side. + By doing so, less pickling needs to be done during operation. Only the task name and data need to be transferred. + Workers can then locally re-construct the task based on the registered tasks. + + Each worker has a state object (state_data). It is provided to each task when it is called. + """ + + def __init__(self, load_function): + super().__init__(RegisterAllTask.register, load_function=load_function) + + @staticmethod + def register(state_data, load_function): + task_dict = dict() + task_dict[DeleteTask.__name__] = DeleteTask + task_dict[ConstructTask.__name__] = ConstructTask + task_dict[ConstructLateTask.__name__] = ConstructLateTask + task_dict[SolveTask.__name__] = SolveTask + task_dict[GetStateTask.__name__] = GetStateTask + task_dict[SetStateTask.__name__] = SetStateTask + task_dict[InitializeTask.__name__] = InitializeTask + task_dict[OutputTask.__name__] = OutputTask + task_dict[ShutdownTask.__name__] = ShutdownTask + state_data["tasks"] = task_dict + state_data["sim_classes"] = dict() + state_data["sim_instances"] = dict() + state_data["load_function"] = load_function + return None + + +def handle_task(state_data, task_descriptor): + name, args, kwargs = task_descriptor + task = state_data["tasks"][name](*args, **kwargs) + # print(f"handling task: {name} args={args} kwargs={kwargs}") + return task(state_data) diff --git a/micro_manager/tasking/worker_main.py b/micro_manager/tasking/worker_main.py new file mode 100644 index 00000000..e9161569 --- /dev/null +++ b/micro_manager/tasking/worker_main.py @@ -0,0 +1,74 @@ +import argparse +import os +from mpi4py import MPI + +# only used for tests that require precice +# pyprecice does not exist in CI, thus dummy is provided in test pipeline +# but for that cwd is needed in module PATH +import sys + +sys.path.append(os.getcwd()) + +from task import handle_task + +from connection import Connection, MPIConnection, SocketConnection +from micro_manager.tools.logging_wrapper import Logger + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--backend", required=True, choices=["mpi", "socket"]) + parser.add_argument("--host", help="IP or localhost") + parser.add_argument("--port", type=int, help="Port to open port in micro manager") + args = parser.parse_args() + + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + log = Logger("Worker", rank=rank) + worker_id = rank + + conn, dst_id, src_id = None, 0, 0 + if args.backend == "mpi": + log.log_info(f"Launched Worker with MPI rank: {rank}") + conn = MPIConnection.connect_to_micro_manager(MPI.Comm.Get_parent()) + else: + log.log_info( + f"Launched Worker with Socket rank: {rank}, IP: {args.host} Port: {args.port}" + ) + conn = SocketConnection.connect_to_micro_manager( + worker_id, args.host, args.port + ) + dst_id = src_id = worker_id + log.log_info(f"Worker rank {rank} connected to parent") + state_data = {} + + # register possible tasks + register_task = None + try: + register_task = conn.recv(src_id) + except Exception: + raise RuntimeError("Failed to recv register tasks") + output = register_task(state_data) + try: + conn.send(dst_id, output) + except Exception: + raise RuntimeError("Failed to send register tasks output") + + while True: + task_descriptor = None + try: + task_descriptor = conn.recv(src_id) + except Exception: + break + + output = None + try: + output = handle_task(state_data, task_descriptor) + except Exception: + break + + try: + conn.send(dst_id, output) + except Exception: + break + + conn.close() diff --git a/micro_manager/tools/misc.py b/micro_manager/tools/misc.py index 9b8b3e40..e1a96652 100644 --- a/micro_manager/tools/misc.py +++ b/micro_manager/tools/misc.py @@ -2,7 +2,7 @@ A collection of miscellaneous functions that are used in various parts of the codebase. """ -from typing import Union, Optional +from typing import Union import numpy as np diff --git a/micro_manager/tools/p2p.py b/micro_manager/tools/p2p.py new file mode 100644 index 00000000..dda85fcf --- /dev/null +++ b/micro_manager/tools/p2p.py @@ -0,0 +1,148 @@ +import hashlib +import numpy as np +from mpi4py import MPI + + +def get_ranks_of_sims(global_ids, rank, comm, global_number_of_sims) -> np.ndarray: + """ + Get the ranks of all simulations. + + Parameters + ---------- + global_ids : list + Global ids on local rank. + rank : int + Rank of simulation. + comm : MPI.Comm + MPI communicator. + global_number_of_sims : int + Global number of sims. + + Returns + ------- + ranks_of_sims : np.ndarray + Array of ranks on which simulations exist. + """ + gids_to_rank = dict() + for gid in global_ids: + gids_to_rank[gid] = rank + + ranks_maps_as_list = comm.allgather(gids_to_rank) + + ranks_of_sims = np.zeros(global_number_of_sims, dtype=np.intc) + for ranks_map in ranks_maps_as_list: + for gid, rank in ranks_map.items(): + ranks_of_sims[gid] = rank + + return ranks_of_sims + + +def create_tag(sim_id: int, src_rank: int, dest_rank: int) -> int: + """ + For a given simulations ID, source rank, and destination rank, a unique tag is created. + + Parameters + ---------- + sim_id : int + Global ID of a simulation. + src_rank : int + Rank on which the simulation lives + dest_rank : int + Rank to which data of a simulation is to be sent to. + + Returns + ------- + tag : int + Unique tag. + """ + send_hashtag = hashlib.sha256() + send_hashtag.update((str(src_rank) + str(sim_id) + str(dest_rank)).encode("utf-8")) + tag = int(send_hashtag.hexdigest()[:6], base=16) + return tag + + +def p2p_comm( + global_ids: list[int], + rank: int, + comm: MPI.Comm, + global_number_of_sims: int, + is_sim_on_this_rank: list[bool], + assoc_active_ids: list[int], + data: list, +) -> list: + """ + Handle process to process communication for a given set of associated active IDs and data. + + Parameters + ---------- + global_ids : list[int] + Global ids on local rank. + rank : int + Rank of simulation. + comm : MPI.Comm + MPI communicator. + global_number_of_sims : int + Global number of sims. + is_sim_on_this_rank: list[bool] + Bool flags of whether a simulation given its gid is on this rank or not. + assoc_active_ids : list[int] + Global IDs of active simulations which are not on this rank and are associated to + the inactive simulations on this rank. + data : list + Complete data from which parts are to be sent and received. + + Returns + ------- + recv_reqs : list + List of MPI requests of receive operations. + """ + rank_of_sim = get_ranks_of_sims(global_ids, rank, comm, global_number_of_sims) + + send_map_local = dict() # keys are global IDs, values are rank to send to + send_map = ( + dict() + ) # keys are global IDs of sims to send, values are ranks to send the sims to + recv_map = ( + dict() + ) # keys are global IDs to receive, values are ranks to receive from + + for i in assoc_active_ids: + # Add simulation and its rank to receive map + recv_map[i] = rank_of_sim[i] + # Add simulation and this rank to local sending map + send_map_local[i] = rank + + # Gather information about which sims to send where, from the sending perspective + send_map_list = comm.allgather(send_map_local) + + for d in send_map_list: + for i, r in d.items(): + if is_sim_on_this_rank[i]: + if i in send_map: + send_map[i].append(r) + else: + send_map[i] = [r] + + # Asynchronous send operations + send_reqs = [] + for gid, send_ranks in send_map.items(): + lid = global_ids.index(gid) + for send_rank in send_ranks: + tag = create_tag(gid, rank, send_rank) + req = comm.isend(data[lid], dest=send_rank, tag=tag) + send_reqs.append(req) + + # Asynchronous receive operations + recv_reqs = [] + for gid, recv_rank in recv_map.items(): + tag = create_tag(gid, recv_rank, rank) + bufsize = ( + 1 << 30 + ) # allocate and use a temporary 1 MiB buffer size https://github.com/mpi4py/mpi4py/issues/389 + req = comm.irecv(bufsize, source=recv_rank, tag=tag) + recv_reqs.append(req) + + # Wait for all non-blocking communication to complete + MPI.Request.Waitall(send_reqs) + + return recv_reqs diff --git a/pyproject.toml b/pyproject.toml index dfc9c62f..0c5dfa7d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" name="micro-manager-precice" dynamic = [ "version" ] dependencies = [ - "pyprecice>=3.2", "numpy", "mpi4py", "psutil" + "pyprecice>=3.2", "numpy", "mpi4py", "psutil", "packaging", "scipy" ] requires-python = ">=3.8" authors = [ @@ -41,7 +41,7 @@ Repository = "https://github.com/precice/micro-manager" micro-manager-precice = "micro_manager:main" [tool.setuptools] -packages=["micro_manager", "micro_manager.adaptivity", "micro_manager.snapshot", "micro_manager.tools"] +packages=["micro_manager", "micro_manager.adaptivity", "micro_manager.snapshot", "micro_manager.tools", "micro_manager.tasking"] [tool.setuptools-git-versioning] enabled = true diff --git a/tests/integration/test_unit_cube/micro_dummy.py b/tests/integration/test_unit_cube/micro_dummy.py index f0244c9b..7c59a1a0 100644 --- a/tests/integration/test_unit_cube/micro_dummy.py +++ b/tests/integration/test_unit_cube/micro_dummy.py @@ -53,3 +53,6 @@ def get_state(self): def set_state(self, state): self._state = copy.deepcopy(state) + + def get_global_id(self): + return self._sim_id diff --git a/tests/unit/test_adaptivity_parallel.py b/tests/unit/test_adaptivity_parallel.py index ba5e78b6..fd16182a 100644 --- a/tests/unit/test_adaptivity_parallel.py +++ b/tests/unit/test_adaptivity_parallel.py @@ -4,6 +4,8 @@ import numpy as np from mpi4py import MPI +from micro_manager.tools.p2p import get_ranks_of_sims +from micro_manager.micro_simulation import create_simulation_class from micro_manager.adaptivity.global_adaptivity import GlobalAdaptivityCalculator @@ -21,6 +23,14 @@ def set_state(self, state): def get_state(self): return self._state.copy() + def solve(self, micro_input, dt): + pass + + +class ModelManager: + def get_instance(self, gid, micro_problem_cls): + return micro_problem_cls(gid) + class TestGlobalAdaptivity(TestCase): def setUp(self): @@ -51,6 +61,14 @@ def test_update_inactive_sims_global_adaptivity(self): return_value="L1" ) + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + adaptivity_controller = GlobalAdaptivityCalculator( self._configurator, 5, @@ -59,7 +77,8 @@ def test_update_inactive_sims_global_adaptivity(self): base_logger=MagicMock(), rank=self._rank, comm=self._comm, - micro_problem_cls=MicroSimulation, + micro_problem_cls=sim_cls, + model_manager=ModelManager(), ) adaptivity_controller._is_sim_active = np.array( @@ -78,7 +97,7 @@ def check_for_activation(i, active): dummy_micro_sims = [] for i in global_ids: - dummy_micro_sims.append(MicroSimulation(i)) + dummy_micro_sims.append(sim_cls(i)) adaptivity_controller._update_inactive_sims(dummy_micro_sims) @@ -125,6 +144,14 @@ def test_update_all_active_sims_global_adaptivity(self): return_value="L2rel" ) + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + adaptivity_controller = GlobalAdaptivityCalculator( self._configurator, 5, @@ -133,14 +160,15 @@ def test_update_all_active_sims_global_adaptivity(self): base_logger=MagicMock(), rank=self._rank, comm=self._comm, - micro_problem_cls=MicroSimulation, + micro_problem_cls=sim_cls, + model_manager=ModelManager(), ) adaptivity_controller._adaptivity_data_names = ["data1", "data2"] dummy_micro_sims = [] for i in global_ids: - dummy_micro_sims.append(MicroSimulation(i)) + dummy_micro_sims.append(sim_cls(i)) adaptivity_controller.compute_adaptivity( 0.1, @@ -180,6 +208,14 @@ def test_communicate_micro_output(self): return_value="L1" ) + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + adaptivity_controller = GlobalAdaptivityCalculator( self._configurator, 5, @@ -188,7 +224,8 @@ def test_communicate_micro_output(self): base_logger=MagicMock(), rank=self._rank, comm=self._comm, - micro_problem_cls=MicroSimulation, + micro_problem_cls=sim_cls, + model_manager=ModelManager(), ) adaptivity_controller._is_sim_active = np.array( @@ -219,6 +256,14 @@ def test_get_ranks_of_sims(self): global_ids = [3, 4] expected_ranks_of_sims = [0, 0, 0, 1, 1] + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + adaptivity_controller = GlobalAdaptivityCalculator( self._configurator, 5, @@ -227,9 +272,10 @@ def test_get_ranks_of_sims(self): base_logger=MagicMock(), rank=self._rank, comm=self._comm, - micro_problem_cls=MicroSimulation, + micro_problem_cls=sim_cls, + model_manager=ModelManager(), ) - actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims() + actual_ranks_of_sims = get_ranks_of_sims(global_ids, self._rank, self._comm, 5) self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) diff --git a/tests/unit/test_adaptivity_serial.py b/tests/unit/test_adaptivity_serial.py index f9ef99b7..cee62f58 100644 --- a/tests/unit/test_adaptivity_serial.py +++ b/tests/unit/test_adaptivity_serial.py @@ -25,6 +25,14 @@ def set_state(self, state): def get_state(self): pass + def solve(self, micro_input, dt): + pass + + +class ModelManager: + def get_instance(self, gid, micro_problem_cls): + return micro_problem_cls(gid) + class TestLocalAdaptivity(TestCase): def setUp(self): @@ -94,6 +102,7 @@ def test_update_similarity_dists(self): configurator, nsims=self._number_of_sims, micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), base_logger=MagicMock(), rank=0, ) @@ -146,6 +155,7 @@ def test_update_active_sims(self): rank=0, comm=MagicMock(), micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), ) adaptivity_controller._refine_const = self._refine_const adaptivity_controller._coarse_const = self._coarse_const @@ -182,6 +192,7 @@ def test_adaptivity_norms(self): base_logger=MagicMock(), rank=0, micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), ) fake_data = np.array([[1], [2], [3]]) @@ -263,6 +274,65 @@ def test_adaptivity_norms(self): ) ) + def test_adaptivity_norms_with_zeros_no_warning(self): + """ + Test that L1rel/L2rel must not raise division-by-zero + warning when data contains zeros. + """ + import warnings + + configurator = MagicMock() + configurator.get_adaptivity_similarity_measure = MagicMock(return_value="L2rel") + configurator.get_output_dir = MagicMock(return_value="output_dir") + configurator.get_micro_file_name = MagicMock( + return_value="test_adaptivity_serial" + ) + adaptivity_l2rel = AdaptivityCalculator( + configurator, + nsims=3, + base_logger=MagicMock(), + rank=0, + micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), + ) + + configurator_l1 = MagicMock() + configurator_l1.get_adaptivity_similarity_measure = MagicMock( + return_value="L1rel" + ) + configurator_l1.get_output_dir = MagicMock(return_value="output_dir") + configurator_l1.get_micro_file_name = MagicMock( + return_value="test_adaptivity_serial" + ) + adaptivity_l1rel = AdaptivityCalculator( + configurator_l1, + nsims=3, + base_logger=MagicMock(), + rank=0, + micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), + ) + + # Data with zeros - previously triggered RuntimeWarning: invalid value in true_divide + data_with_zeros = np.array([[0.0], [0.0], [1.0]]) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("error", RuntimeWarning) + result_l2rel = adaptivity_l2rel._l2rel(data_with_zeros) + self.assertEqual( + len(w), 0, "L2rel must not raise RuntimeWarning with zero data" + ) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("error", RuntimeWarning) + result_l1rel = adaptivity_l1rel._l1rel(data_with_zeros) + self.assertEqual( + len(w), 0, "L1rel must not raise RuntimeWarning with zero data" + ) + + # When both are 0, relative diff should be 0 (since numerator is 0) + self.assertEqual(result_l2rel[0, 1], 0.0) + self.assertEqual(result_l1rel[0, 1], 0.0) + def test_associate_active_to_inactive(self): """ Test functionality to associate inactive sims to active ones, in the class AdaptivityCalculator. @@ -280,6 +350,7 @@ def test_associate_active_to_inactive(self): base_logger=MagicMock(), rank=0, micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), ) adaptivity_controller._refine_const = self._refine_const adaptivity_controller._coarse_const = self._coarse_const @@ -323,6 +394,7 @@ def test_update_inactive_sims_local_adaptivity(self): rank=0, comm=MPI.COMM_WORLD, micro_problem_cls=MicroSimulation, + model_manager=ModelManager(), ) adaptivity_controller._refine_const = self._refine_const adaptivity_controller._coarse_const = self._coarse_const diff --git a/tests/unit/test_domain_decomposition.py b/tests/unit/test_domain_decomposition.py index aeba3d14..a3d1e14b 100644 --- a/tests/unit/test_domain_decomposition.py +++ b/tests/unit/test_domain_decomposition.py @@ -1,4 +1,5 @@ -from unittest import TestCase +from unittest import TestCase, skip +from unittest.mock import MagicMock import numpy as np from micro_manager.domain_decomposition import DomainDecomposer @@ -21,18 +22,20 @@ def setUp(self) -> None: 2, ] + self._configuration_mock = MagicMock() + def test_rank2_out_of_4_2d(self): """ Check bounds for rank 2 in a setting of axis-wise ranks: [2, 2] """ - rank = 2 - size = 4 - ranks_per_axis = [2, 2] - domain_decomposer = DomainDecomposer(rank, size) - domain_decomposer._dims = 2 - mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds_2d, ranks_per_axis + self._configuration_mock.get_decomposition_type.return_value = "uniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_2d ) + self._configuration_mock.get_ranks_per_axis.return_value = [2, 2] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=2, size=4) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() self.assertTrue(np.allclose(mesh_bounds, [0, 0.5, 1, 2])) @@ -40,14 +43,14 @@ def test_rank1_out_of_4_3d(self): """ Check bounds for rank 1 in a setting of axis-wise ranks: [2, 2, 1] """ - rank = 1 - size = 4 - ranks_per_axis = [2, 2, 1] - domain_decomposer = DomainDecomposer(rank, size) - domain_decomposer._dims = 3 - mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds_3d, ranks_per_axis + self._configuration_mock.get_decomposition_type.return_value = "uniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_3d ) + self._configuration_mock.get_ranks_per_axis.return_value = [2, 2, 1] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=1, size=4) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() self.assertTrue(np.allclose(mesh_bounds, [0.0, 1, -2, 0.0, -2, 8])) @@ -55,14 +58,14 @@ def test_rank5_outof_10_3d(self): """ Test domain decomposition for rank 5 in a setting of axis-wise ranks: [1, 2, 5] """ - rank = 5 - size = 10 - ranks_per_axis = [1, 2, 5] - domain_decomposer = DomainDecomposer(rank, size) - domain_decomposer._dims = 3 - mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds_3d, ranks_per_axis + self._configuration_mock.get_decomposition_type.return_value = "uniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_3d ) + self._configuration_mock.get_ranks_per_axis.return_value = [1, 2, 5] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=5, size=10) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() self.assertTrue(np.allclose(mesh_bounds, [-1, 1, 0, 2, 2, 4])) @@ -70,14 +73,14 @@ def test_rank10_out_of_32_3d(self): """ Test domain decomposition for rank 10 in a setting of axis-wise ranks: [4, 1, 8] """ - rank = 10 - size = 32 - ranks_per_axis = [4, 1, 8] - domain_decomposer = DomainDecomposer(rank, size) - domain_decomposer._dims = 3 - mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds_3d, ranks_per_axis + self._configuration_mock.get_decomposition_type.return_value = "uniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_3d ) + self._configuration_mock.get_ranks_per_axis.return_value = [4, 1, 8] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=10, size=32) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() self.assertTrue(np.allclose(mesh_bounds, [0, 0.5, -2, 2, 0.5, 1.75])) @@ -85,13 +88,203 @@ def test_rank7_out_of_16_3d(self): """ Test domain decomposition for rank 7 in a setting of axis-wise ranks: [8, 2, 1] """ - rank = 7 - size = 16 - ranks_per_axis = [8, 2, 1] - domain_decomposer = DomainDecomposer(rank, size) - domain_decomposer._dims = 3 - mesh_bounds = domain_decomposer.get_local_mesh_bounds( - self._macro_bounds_3d, ranks_per_axis + self._configuration_mock.get_decomposition_type.return_value = "uniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_3d ) + self._configuration_mock.get_ranks_per_axis.return_value = [8, 2, 1] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=7, size=16) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() self.assertTrue(np.allclose(mesh_bounds, [0.75, 1, -2, 0, -2, 8])) + + +class TestNonUniformDomainDecomposition(TestCase): + def setUp(self) -> None: + self._macro_bounds_3d = [ + -1, + 1, + -2, + 2, + -2, + 8, + ] + self._macro_bounds_2d = [ + 0, + 1, + 0, + 2, + ] + + self._configuration_mock = MagicMock() + + def test_nonuniform_rank2_out_of_4_2d(self): + """ + Check non-uniform bounds for rank 2 in a setting of axis-wise ranks: [2, 2]. + Along each axis, the local width doubles from one rank to the next. + """ + self._configuration_mock.get_decomposition_type.return_value = "nonuniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_2d + ) + self._configuration_mock.get_ranks_per_axis.return_value = [2, 2] + self._configuration_mock.get_minimum_access_region_size.return_value = [] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=2, size=4) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() + + self.assertTrue(np.allclose(mesh_bounds, [0.0, 1.0 / 3.0, 2.0 / 3.0, 2.0])) + + def test_nonuniform_rank15_out_of_128_2d(self): + """ + Check non-uniform bounds for rank 15 in a setting of axis-wise ranks: [16, 8]. + Along each axis, the local width doubles from one rank to the next. + """ + self._configuration_mock.get_decomposition_type.return_value = "nonuniform" + self._configuration_mock.get_macro_domain_bounds.return_value = [0, 1, 0, 0.5] + self._configuration_mock.get_ranks_per_axis.return_value = [16, 8] + self._configuration_mock.get_minimum_access_region_size.return_value = [ + 1.0 / 256.0, + 1.0 / 128.0, + ] + + # In a 16 x 8 grid, rank 15 is in the lower right corner. + domain_decomposer = DomainDecomposer( + self._configuration_mock, rank=15, size=128 + ) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() + + self.assertTrue( + np.allclose(mesh_bounds, [0.756153, 1.0, 0.0, 0.019664], atol=1e-5) + ) + + # In a 16 x 8 grid, rank 112 is in the lower right corner. + domain_decomposer = DomainDecomposer( + self._configuration_mock, rank=112, size=128 + ) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() + + self.assertTrue( + np.allclose(mesh_bounds, [0.0, 1.0 / 256.0, 0.364631, 0.5], atol=1e-5) + ) + + def test_nonuniform_rank1_out_of_4_3d(self): + """ + Check non-uniform bounds for rank 1 in a setting of axis-wise ranks: [2, 2, 1]. + """ + self._configuration_mock.get_decomposition_type.return_value = "nonuniform" + self._configuration_mock.get_macro_domain_bounds.return_value = ( + self._macro_bounds_3d + ) + self._configuration_mock.get_ranks_per_axis.return_value = [2, 2, 1] + self._configuration_mock.get_minimum_access_region_size.return_value = [] + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=1, size=4) + mesh_bounds = domain_decomposer.get_local_mesh_bounds() + + self.assertTrue( + np.allclose(mesh_bounds, [-1.0 / 3.0, 1.0, -2.0, -2.0 / 3.0, -2.0, 8.0]) + ) + + def test_nonuniform_invalid_processor_count_raises(self): + """ + A mismatch between `ranks_per_axis` and communicator size should raise a ValueError. + """ + self._configuration_mock.get_decomposition_type.return_value = "nonuniform" + + domain_decomposer = DomainDecomposer(self._configuration_mock, rank=0, size=4) + + with self.assertRaises(ValueError): + domain_decomposer.get_local_mesh_bounds() + + +class TestDuplicateCoordFiltering(TestCase): + """ + Test that duplicate vertex coordinates returned by preCICE on rank boundaries + are correctly filtered, with lower-ranked ranks taking ownership. + """ + + def setUp(self) -> None: + self._configuration_mock = MagicMock() + + def test_no_duplicates(self): + """ + If there are no shared coords across ranks, nothing should be filtered. + """ + all_coords = [ + np.array([[0.0, 0.0], [0.5, 0.0]]), + np.array([[1.0, 0.0], [1.5, 0.0]]), + ] + all_ids = [[0, 1], [2, 3]] + + self._configuration_mock.get_decomposition_type.return_value = "uniform" + + coords, ids = DomainDecomposer( + self._configuration_mock, rank=0, size=2 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords), 2) + + coords, ids = DomainDecomposer( + self._configuration_mock, rank=1, size=2 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords), 2) + + def test_duplicate_on_boundary_rank0_keeps(self): + """ + A coord shared between rank 0 and rank 1 should be kept by rank 0 + and dropped by rank 1. + """ + shared = [0.5, 0.0] + all_coords = [ + np.array([[0.0, 0.0], shared]), + np.array([shared, [1.0, 0.0]]), + ] + all_ids = [[0, 1], [1, 2]] + + self._configuration_mock.get_decomposition_type.return_value = "uniform" + + # Rank 0 should keep both its coords + coords0, ids0 = DomainDecomposer( + self._configuration_mock, rank=0, size=2 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords0), 2) + self.assertTrue(np.allclose(coords0[1], shared)) + + # Rank 1 should drop the shared coord + coords1, ids1 = DomainDecomposer( + self._configuration_mock, rank=1, size=2 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords1), 1) + self.assertTrue(np.allclose(coords1[0], [1.0, 0.0])) + + def test_duplicate_on_boundary_three_ranks(self): + """ + A coord shared across three ranks should only be kept by rank 0. + """ + shared = [0.5, 0.5] + all_coords = [ + np.array([shared, [0.0, 0.0]]), + np.array([shared, [1.0, 0.0]]), + np.array([shared, [2.0, 0.0]]), + ] + all_ids = [[0, 1], [0, 2], [0, 3]] + + self._configuration_mock.get_decomposition_type.return_value = "uniform" + + coords0, _ = DomainDecomposer( + self._configuration_mock, rank=0, size=3 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords0), 2) + + coords1, _ = DomainDecomposer( + self._configuration_mock, rank=1, size=3 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords1), 1) + self.assertTrue(np.allclose(coords1[0], [1.0, 0.0])) + + coords2, _ = DomainDecomposer( + self._configuration_mock, rank=2, size=3 + ).filter_duplicate_coords(all_coords, all_ids) + self.assertEqual(len(coords2), 1) + self.assertTrue(np.allclose(coords2[0], [2.0, 0.0])) diff --git a/tests/unit/test_global_adaptivity_lb.py b/tests/unit/test_global_adaptivity_lb.py deleted file mode 100644 index d60ae12b..00000000 --- a/tests/unit/test_global_adaptivity_lb.py +++ /dev/null @@ -1,295 +0,0 @@ -import unittest -from unittest import TestCase -from unittest.mock import MagicMock - -import numpy as np -from mpi4py import MPI - -from micro_manager.adaptivity.global_adaptivity_lb import GlobalAdaptivityLBCalculator - - -class MicroSimulation: - def __init__(self, global_id) -> None: - self._global_id = global_id - self._state = [global_id] * 3 - - def get_global_id(self): - return self._global_id - - def set_state(self, state): - self._state = state - - def get_state(self): - return self._state.copy() - - -class TestGlobalAdaptivityLB(TestCase): - def setUp(self): - self._comm = MPI.COMM_WORLD - self._rank = self._comm.Get_rank() - self._size = self._comm.Get_size() - - self._configurator = MagicMock() - self._configurator.get_micro_file_name = MagicMock( - return_value="test_global_adaptivity_lb" - ) - self._configurator.get_adaptivity_similarity_measure = MagicMock( - return_value="L1" - ) - self._configurator.get_output_dir = MagicMock(return_value="output_dir") - - self._configurator.get_load_balancing_threshold = MagicMock(return_value=0) - - @unittest.skipUnless( - MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." - ) - def test_redistribute_active_sims_two_ranks(self): - """ - Test load balancing functionality to redistribute active simulations. - Run this test in parallel using MPI with 2 ranks. - """ - global_number_of_sims = 8 - - if self._rank == 0: - global_ids = [0, 1, 2, 3] - expected_global_ids = [2, 3] - elif self._rank == 1: - global_ids = [4, 5, 6, 7] - expected_global_ids = [4, 5, 6, 7, 0, 1] - - expected_ranks_of_sims = [1, 1, 0, 0, 1, 1, 1, 1] - - adaptivity_controller = GlobalAdaptivityLBCalculator( - self._configurator, - global_number_of_sims, - global_ids, - participant=MagicMock(), - base_logger=MagicMock(), - rank=self._rank, - comm=self._comm, - micro_problem_cls=MicroSimulation, - ) - - adaptivity_controller._is_sim_active = np.array( - [True, True, True, True, False, False, False, False] - ) - - micro_sims = [] - for i in global_ids: - micro_sims.append(MicroSimulation(i)) - - adaptivity_controller._redistribute_active_sims(micro_sims) - - actual_global_ids = [] - for sim in micro_sims: - actual_global_ids.append(sim.get_global_id()) - - self.assertEqual(actual_global_ids, expected_global_ids) - - actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims() - - self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) - - @unittest.skipUnless( - MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." - ) - def test_redistribute_inactive_sims_two_ranks(self): - """ - Test load balancing functionality to redistribute inactive simulations. - Run this test in parallel using MPI with 2 ranks. - """ - global_number_of_sims = 5 - - if self._rank == 0: - global_ids = [0, 2] - expected_global_ids = [0, 2, 4] - elif self._rank == 1: - global_ids = [1, 3, 4] - expected_global_ids = [1, 3] - - expected_ranks_of_sims = [0, 1, 0, 1, 0] - - adaptivity_controller = GlobalAdaptivityLBCalculator( - self._configurator, - global_number_of_sims, - global_ids, - participant=MagicMock(), - base_logger=MagicMock(), - rank=self._rank, - comm=self._comm, - micro_problem_cls=MicroSimulation, - ) - - adaptivity_controller._is_sim_active = np.array( - [True, True, False, False, False] - ) - adaptivity_controller._sim_is_associated_to = [-2, -2, 0, 1, 0] - - micro_sims = [] - for i in global_ids: - micro_sims.append(MicroSimulation(i)) - - adaptivity_controller._redistribute_inactive_sims(micro_sims) - - self.assertEqual(global_ids, expected_global_ids) - - actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims() - - self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) - - @unittest.skipUnless( - MPI.COMM_WORLD.Get_size() == 4, "This test only works with 4 ranks." - ) - def test_redistribute_active_sims_four_ranks(self): - """ - Test load balancing functionality to redistribute active simulations. The load balancing is one in two steps. - Run this test in parallel using MPI with 4 ranks. - """ - global_number_of_sims = 15 - - if self._rank == 0: - global_ids = [0, 1, 2, 3] - expected_global_ids = [1, 2, 3] - elif self._rank == 1: - global_ids = [4, 5, 6, 7] - expected_global_ids = [4, 5, 6, 7, 0] - elif self._rank == 2: - global_ids = [8, 9, 10, 11] - expected_global_ids = [8, 9, 10, 11, 12] - elif self._rank == 3: - global_ids = [12, 13, 14] - expected_global_ids = [13, 14] - - expected_ranks_of_sims = [1, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3] - - adaptivity_controller = GlobalAdaptivityLBCalculator( - self._configurator, - global_number_of_sims, - global_ids, - participant=MagicMock(), - base_logger=MagicMock(), - rank=self._rank, - comm=self._comm, - micro_problem_cls=MicroSimulation, - ) - - adaptivity_controller._is_sim_active = np.array( - [ - True, - True, - True, - False, - False, - False, - False, - False, - True, - False, - False, - False, - True, - True, - True, - ] - ) - - micro_sims = [] - for i in global_ids: - micro_sims.append(MicroSimulation(i)) - - adaptivity_controller._redistribute_active_sims(micro_sims) - - actual_global_ids = [] - for sim in micro_sims: - actual_global_ids.append(sim.get_global_id()) - - self.assertEqual(actual_global_ids, expected_global_ids) - - actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims() - - self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) - - @unittest.skipUnless( - MPI.COMM_WORLD.Get_size() == 4, "This test only works with 4 ranks." - ) - def test_redistribute_inactive_sims_four_ranks(self): - """ - Test load balancing functionality to redistribute inactive simulations. - Run this test in parallel using MPI with 4 ranks. - """ - global_number_of_sims = 15 - - if self._rank == 0: - global_ids = [1, 2, 3] - expected_global_ids = [1, 2, 4, 9, 10] - elif self._rank == 1: - global_ids = [4, 5, 6, 7, 12] - expected_global_ids = [6, 7, 12] - elif self._rank == 2: - global_ids = [0, 8, 9, 10, 11] - expected_global_ids = [0, 8, 11, 3] - elif self._rank == 3: - global_ids = [13, 14] - expected_global_ids = [13, 14, 5] - - expected_ranks_of_sims = [2, 0, 0, 2, 0, 3, 1, 1, 2, 0, 0, 2, 1, 3, 3] - - adaptivity_controller = GlobalAdaptivityLBCalculator( - self._configurator, - global_number_of_sims, - global_ids, - participant=MagicMock(), - base_logger=MagicMock(), - rank=self._rank, - comm=self._comm, - micro_problem_cls=MicroSimulation, - ) - - adaptivity_controller._is_sim_active = np.array( - [ - True, - True, - True, - False, - False, - False, - False, - False, - True, - False, - False, - False, - True, - True, - True, - ] - ) - adaptivity_controller._sim_is_associated_to = [ - -2, - -2, - -2, - 0, - 1, - 13, - 12, - 12, - -2, - 1, - 2, - 8, - -2, - -2, - -2, - ] - - micro_sims = [] - for i in global_ids: - micro_sims.append(MicroSimulation(i)) - - adaptivity_controller._redistribute_inactive_sims(micro_sims) - - self.assertEqual(global_ids, expected_global_ids) - - actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims() - - self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) diff --git a/tests/unit/test_interpolation.py b/tests/unit/test_interpolation.py index 8f5cb9e6..47062bba 100644 --- a/tests/unit/test_interpolation.py +++ b/tests/unit/test_interpolation.py @@ -48,3 +48,29 @@ def test_nearest_neighbor(self): self.assertListEqual( nearest_neighbor_index.tolist(), expected_nearest_neighbor_index ) + + def test_interpolation_exact_point(self): + """ + Test that if interpolation point exactly matches a neighbor, that value is returned. + """ + coords = [[0, 0, 0], [1, 0, 0], [2, 0, 0]] + point = [1, 0, 0] + values = [10, 20, 30] + + interpolation = Interpolation(MagicMock()) + result = interpolation.interpolate(coords, point, values) + self.assertEqual(result, 20) + + def test_nearest_neighbor_k_larger_than_coords(self): + """ + Test that k is reset when larger than number of available neighbors. + """ + coords = [[0, 0, 0], [1, 0, 0]] + inter_point = [0.5, 0, 0] + k = 5 # larger than len(coords) + + mock_logger = MagicMock() + interpolation = Interpolation(mock_logger) + indices = interpolation.get_nearest_neighbor_indices(coords, inter_point, k) + self.assertEqual(len(indices), 2) + mock_logger.log_info.assert_called_once() diff --git a/tests/unit/test_load_balancing.py b/tests/unit/test_load_balancing.py new file mode 100644 index 00000000..e294c261 --- /dev/null +++ b/tests/unit/test_load_balancing.py @@ -0,0 +1,566 @@ +import unittest +from unittest import TestCase +from unittest.mock import MagicMock + +import numpy as np +from mpi4py import MPI + +from micro_manager.load_balancing import LoadBalancer, ActiveBalancer +from micro_manager.micro_simulation import create_simulation_class +from micro_manager.tools.p2p import get_ranks_of_sims + + +class MicroSimulation: + def __init__(self, global_id) -> None: + self._global_id = global_id + self._state = [global_id] * 3 + + def get_global_id(self): + return self._global_id + + def set_state(self, state): + self._state = state + + def get_state(self): + return self._state.copy() + + def solve(self, micro_input, dt): + pass + + +class ModelManager: + def __init__(self, cls): + self._cls = cls + + def get_instance(self, gid, micro_problem_cls): + assert micro_problem_cls == self._cls + return micro_problem_cls(gid) + + def is_stateless(self, _): + return True + + def get_instance_by_name(self, gid: int, name: str, *, late_init: bool = False): + return self._cls(gid, late_init=late_init) + + +class TestLBTime(TestCase): + def setUp(self): + self._comm = MPI.COMM_WORLD + self._rank = self._comm.Get_rank() + self._size = self._comm.Get_size() + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." + ) + def test_redistribute_no_time_two_ranks(self): + """ + Test load balancing functionality to redistribute active simulations. + Run this test in parallel using MPI with 2 ranks. + """ + global_number_of_sims = 8 + + # assume lpt partitioning + if self._rank == 0: + global_ids = [2, 3] + expected_global_ids = [1, 3, 5, 7] + else: + global_ids = [0, 1, 4, 5, 6, 7] + expected_global_ids = [0, 2, 4, 6] + expected_ranks_of_sims = [1, 0, 1, 0, 1, 0, 1, 0] + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + micro_sims.append(sim_cls(i)) + + config = MagicMock() + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + load_balancer = LoadBalancer( + MagicMock(), + ModelManager(sim_cls), + None, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + + load_balancer.balance() + + actual_global_ids = [] + for sim in micro_sims: + actual_global_ids.append(sim.get_global_id()) + self.assertListEqual(sorted(actual_global_ids), expected_global_ids) + + actual_ranks_of_sims = get_ranks_of_sims( + global_ids, self._rank, self._comm, global_number_of_sims + ) + self.assertListEqual(list(expected_ranks_of_sims), list(actual_ranks_of_sims)) + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." + ) + def test_redistribute_no_time_with_inactive_sims_two_ranks(self): + """ + Test load balancing functionality to redistribute inactive simulations. + Run this test in parallel using MPI with 2 ranks. + """ + global_number_of_sims = 8 + + class GlobalAdaptivityCalculator: + def __init__(self, active_ids): + self._active_ids = active_ids + + def get_active_sim_local_ids(self): + return self._active_ids + + def set_is_on_rank(self, *args, **kwargs): + pass + + # assume lpt partitioning + if self._rank == 0: + global_ids = [2, 3] + expected_global_ids = [1, 3, 5, 7] + adaptivity = GlobalAdaptivityCalculator([1]) + else: + global_ids = [0, 1, 4, 5, 6, 7] + expected_global_ids = [0, 2, 4, 6] + adaptivity = GlobalAdaptivityCalculator([1, 2, 3, 4, 5]) + expected_ranks_of_sims = [1, 0, 1, 0, 1, 0, 1, 0] + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + if i in [0, 2]: + micro_sims.append(None) + else: + micro_sims.append(sim_cls(i)) + + config = MagicMock() + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + load_balancer = LoadBalancer( + MagicMock(), + ModelManager(sim_cls), + adaptivity, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + + load_balancer.balance() + + self.assertListEqual(sorted(global_ids), expected_global_ids) + + actual_ranks_of_sims = get_ranks_of_sims( + global_ids, self._rank, self._comm, global_number_of_sims + ) + self.assertListEqual(list(expected_ranks_of_sims), list(actual_ranks_of_sims)) + + +class GlobalAdaptivityCalculator: + def __init__(self, is_active, active_lids, active_gids, associated, on_rank): + self._is_sim_active = is_active + self._active_lids = active_lids + self._active_gids = active_gids + self._sim_is_associated_to = associated + self._on_rank = on_rank + + def get_active_sim_local_ids(self): + return self._active_lids + + def get_active_sim_global_ids(self): + return self._active_gids + + def set_is_on_rank(self, *args, **kwargs): + pass + + +class TestLBActive(TestCase): + def setUp(self): + self._comm = MPI.COMM_WORLD + self._rank = self._comm.Get_rank() + self._size = self._comm.Get_size() + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." + ) + def test_redistribute_active_sims_two_ranks(self): + """ + Test load balancing functionality to redistribute active simulations. + Run this test in parallel using MPI with 2 ranks. + """ + config = MagicMock() + config.get_load_balancing_threshold = MagicMock(return_value=0) + config.turn_on_load_balancing_inactive = MagicMock(return_value=False) + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + + global_number_of_sims = 8 + + if self._rank == 0: + global_ids = [0, 1, 2, 3] + expected_global_ids = [2, 3] + active_lids = [0, 1, 2, 3] + active_gids = [0, 1, 2, 3] + elif self._rank == 1: + global_ids = [4, 5, 6, 7] + expected_global_ids = [4, 5, 6, 7, 0, 1] + active_lids = [] + active_gids = [] + expected_ranks_of_sims = [1, 1, 0, 0, 1, 1, 1, 1] + + adaptivity_controller = GlobalAdaptivityCalculator( + np.array([True, True, True, True, False, False, False, False]), + active_lids, + active_gids, + [-2] * global_number_of_sims, + [gid in global_ids for gid in range(global_number_of_sims)], + ) + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + micro_sims.append(sim_cls(i)) + + load_balancer = ActiveBalancer( + MagicMock(), + ModelManager(sim_cls), + adaptivity_controller, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + + load_balancer.balance() + + actual_global_ids = [] + for sim in micro_sims: + actual_global_ids.append(sim.get_global_id()) + self.assertListEqual(actual_global_ids, expected_global_ids) + + actual_ranks_of_sims = get_ranks_of_sims( + actual_global_ids, self._rank, self._comm, global_number_of_sims + ) + self.assertListEqual(expected_ranks_of_sims, list(actual_ranks_of_sims)) + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 2, "This test only works with 2 ranks." + ) + def test_redistribute_inactive_sims_two_ranks(self): + """ + Test load balancing functionality to redistribute inactive simulations. + Run this test in parallel using MPI with 2 ranks. + """ + config = MagicMock() + config.get_load_balancing_threshold = MagicMock(return_value=0) + config.turn_on_load_balancing_inactive = MagicMock(return_value=True) + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + global_number_of_sims = 5 + + if self._rank == 0: + global_ids = [0, 2] + expected_global_ids = [0, 2, 4] + active_lids = [0] + active_gids = [0] + elif self._rank == 1: + global_ids = [1, 3, 4] + expected_global_ids = [1, 3] + active_lids = [0] + active_gids = [1] + expected_ranks_of_sims = [0, 1, 0, 1, 0] + + adaptivity_controller = GlobalAdaptivityCalculator( + np.array([True, True, False, False, False]), + active_lids, + active_gids, + [-2, -2, 0, 1, 0], + [gid in global_ids for gid in range(global_number_of_sims)], + ) + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + micro_sims.append(sim_cls(i)) + + load_balancer = ActiveBalancer( + MagicMock(), + ModelManager(sim_cls), + adaptivity_controller, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + load_balancer._bypass_skip = True + load_balancer._bypass_active = True + + load_balancer.balance() + self.assertEqual(global_ids, expected_global_ids) + + self.assertListEqual(global_ids, expected_global_ids) + actual_ranks_of_sims = get_ranks_of_sims( + global_ids, self._rank, self._comm, global_number_of_sims + ) + + self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 4, "This test only works with 4 ranks." + ) + def test_redistribute_active_sims_four_ranks(self): + """ + Test load balancing functionality to redistribute active simulations. The load balancing is one in two steps. + Run this test in parallel using MPI with 4 ranks. + """ + config = MagicMock() + config.get_load_balancing_threshold = MagicMock(return_value=0) + config.turn_on_load_balancing_inactive = MagicMock(return_value=False) + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + + global_number_of_sims = 15 + + # active_gids_global = [0, 1, 2, 8, 12, 13, 14] + if self._rank == 0: + global_ids = [0, 1, 2, 3] + expected_global_ids = [1, 2, 3] + active_lids = [0, 1, 2] + active_gids = [0, 1, 2] + elif self._rank == 1: + global_ids = [4, 5, 6, 7] + expected_global_ids = [4, 5, 6, 7, 0] + active_lids = [] + active_gids = [] + elif self._rank == 2: + global_ids = [8, 9, 10, 11] + expected_global_ids = [8, 9, 10, 11, 12] + active_lids = [0] + active_gids = [8] + elif self._rank == 3: + global_ids = [12, 13, 14] + expected_global_ids = [13, 14] + active_lids = [0, 1, 2] + active_gids = [12, 13, 14] + expected_ranks_of_sims = [1, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3] + + adaptivity_controller = GlobalAdaptivityCalculator( + np.array( + [ + True, + True, + True, + False, + False, + False, + False, + False, + True, + False, + False, + False, + True, + True, + True, + ] + ), + active_lids, + active_gids, + [-2] * global_number_of_sims, + [gid in global_ids for gid in range(global_number_of_sims)], + ) + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + micro_sims.append(sim_cls(i)) + + load_balancer = ActiveBalancer( + MagicMock(), + ModelManager(sim_cls), + adaptivity_controller, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + + load_balancer.balance() + + actual_global_ids = [] + for sim in micro_sims: + actual_global_ids.append(sim.get_global_id()) + self.assertListEqual(actual_global_ids, expected_global_ids) + + actual_ranks_of_sims = get_ranks_of_sims( + actual_global_ids, self._rank, self._comm, global_number_of_sims + ) + self.assertListEqual(expected_ranks_of_sims, list(actual_ranks_of_sims)) + + @unittest.skipUnless( + MPI.COMM_WORLD.Get_size() == 4, "This test only works with 4 ranks." + ) + def test_redistribute_inactive_sims_four_ranks(self): + """ + Test load balancing functionality to redistribute inactive simulations. + Run this test in parallel using MPI with 4 ranks. + """ + config = MagicMock() + config.get_load_balancing_threshold = MagicMock(return_value=0) + config.turn_on_load_balancing_inactive = MagicMock(return_value=True) + config.turn_on_load_balancing = MagicMock(return_value=True) + config.get_load_balancing_partitioning = MagicMock(return_value="lpt") + global_number_of_sims = 15 + + # global_active_gids = [0, 1, 2, 8, 12, 13, 14] + if self._rank == 0: + global_ids = [1, 2, 3] + expected_global_ids = [1, 2, 4, 9, 10] + active_lids = [0, 1] + active_gids = [1, 2] + elif self._rank == 1: + global_ids = [4, 5, 6, 7, 12] + expected_global_ids = [6, 7, 12] + active_lids = [4] + active_gids = [12] + elif self._rank == 2: + global_ids = [0, 8, 9, 10, 11] + expected_global_ids = [0, 8, 11, 3] + active_lids = [0, 1] + active_gids = [0, 8] + elif self._rank == 3: + global_ids = [13, 14] + expected_global_ids = [13, 14, 5] + active_lids = [0, 1] + active_gids = [13, 14] + expected_ranks_of_sims = [2, 0, 0, 2, 0, 3, 1, 1, 2, 0, 0, 2, 1, 3, 3] + + adaptivity_controller = GlobalAdaptivityCalculator( + np.array( + [ + True, + True, + True, + False, + False, + False, + False, + False, + True, + False, + False, + False, + True, + True, + True, + ] + ), + active_lids, + active_gids, + [-2, -2, -2, 0, 1, 13, 12, 12, -2, 1, 2, 8, -2, -2, -2], + [gid in global_ids for gid in range(global_number_of_sims)], + ) + + sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + __file__, + 1, + None, + ) + + micro_sims = [] + for i in global_ids: + micro_sims.append(sim_cls(i)) + + load_balancer = ActiveBalancer( + MagicMock(), + ModelManager(sim_cls), + adaptivity_controller, + lambda sim: sim.get_state(), + lambda sim, state: sim.set_state(state), + MagicMock(), + config, + micro_sims, + global_ids, + global_number_of_sims, + self._comm, + self._rank, + ) + load_balancer._bypass_skip = True + load_balancer._bypass_active = True + + load_balancer.balance() + + self.assertEqual(global_ids, expected_global_ids) + + self.assertListEqual(global_ids, expected_global_ids) + actual_ranks_of_sims = get_ranks_of_sims( + global_ids, self._rank, self._comm, global_number_of_sims + ) + + self.assertTrue(np.array_equal(expected_ranks_of_sims, actual_ranks_of_sims)) diff --git a/tests/unit/test_micro_manager.py b/tests/unit/test_micro_manager.py index 189c9303..35ae8ea4 100644 --- a/tests/unit/test_micro_manager.py +++ b/tests/unit/test_micro_manager.py @@ -21,6 +21,15 @@ def solve(self, macro_data, dt): "Micro-Vector-Data": macro_data["Macro-Vector-Data"] + 1, } + def get_global_id(self): + pass + + def get_state(self): + return None + + def set_state(self, state): + pass + class TestFunctioncalls(TestCase): def setUp(self): @@ -50,7 +59,6 @@ def test_micromanager_constructor(self): """ manager = micro_manager.MicroManagerCoupling("micro-manager-config.json") - self.assertListEqual(manager._macro_bounds, self.macro_bounds) self.assertListEqual(manager._read_data_names, self.fake_read_data_names) self.assertListEqual(self.fake_write_data_names, manager._write_data_names) self.assertEqual(manager._micro_n_out, 10) @@ -64,7 +72,6 @@ def test_initialization(self): self.assertEqual(manager._micro_dt, 0.1) # from Interface.initialize self.assertEqual(manager._global_number_of_sims, 4) - self.assertListEqual(manager._macro_bounds, self.macro_bounds) self.assertListEqual(manager._mesh_vertex_ids.tolist(), [0, 1, 2, 3]) self.assertEqual(len(manager._micro_sims), 4) self.assertEqual( diff --git a/tests/unit/test_micro_simulation.py b/tests/unit/test_micro_simulation.py new file mode 100644 index 00000000..b0d2fbf8 --- /dev/null +++ b/tests/unit/test_micro_simulation.py @@ -0,0 +1,376 @@ +""" +Tests for micro_simulation.py covering MicroSimulationInterface, +MicroSimulationLocal, MicroSimulationClass, and create_simulation_class. +""" + +import unittest +import warnings +from unittest.mock import MagicMock + +from micro_manager.micro_simulation import ( + MicroSimulationInterface, + MicroSimulationLocal, + create_simulation_class, +) + + +class MinimalSim(MicroSimulationInterface): + """Minimal implementation of MicroSimulationInterface.""" + + def __init__(self, gid): + self._gid = gid + + def solve(self, micro_sim_input, dt): + return {"out": 1} + + def get_state(self): + return self._gid + + def set_state(self, state): + self._gid = state + + def get_global_id(self): + return self._gid + + def set_global_id(self, gid): + self._gid = gid + + +class SimWithInitialize(MinimalSim): + def initialize(self, data=None): + return {"init": True} + + +class SimWithOutput(MinimalSim): + def output(self): + pass + + +class TestMicroSimulationInterface(unittest.TestCase): + def test_requires_initialize_false_when_not_overridden(self): + sim = MinimalSim(0) + self.assertFalse(sim.requires_initialize()) + + def test_requires_initialize_true_when_overridden(self): + sim = SimWithInitialize(0) + self.assertTrue(sim.requires_initialize()) + + def test_requires_output_false_when_not_overridden(self): + sim = MinimalSim(0) + self.assertFalse(sim.requires_output()) + + def test_requires_output_true_when_overridden(self): + sim = SimWithOutput(0) + self.assertTrue(sim.requires_output()) + + def test_default_initialize_returns_none(self): + sim = MinimalSim(0) + self.assertIsNone(sim.initialize()) + + def test_default_output_returns_none(self): + sim = MinimalSim(0) + self.assertIsNone(sim.output()) + + +class TestMicroSimulationLocal(unittest.TestCase): + def test_solve(self): + local = MicroSimulationLocal(0, False, MinimalSim) + result = local.solve({"in": 1}, 0.1) + self.assertEqual(result, {"out": 1}) + + def test_get_set_state(self): + local = MicroSimulationLocal(0, False, MinimalSim) + local.set_state(42) + self.assertEqual(local.get_state(), 42) + + def test_get_set_global_id(self): + local = MicroSimulationLocal(5, False, MinimalSim) + self.assertEqual(local.get_global_id(), 5) + local.set_global_id(99) + self.assertEqual(local.get_global_id(), 99) + + def test_late_init_sets_instance_gid_to_minus_one(self): + """When late_init=True, the wrapped instance should be constructed with gid=-1.""" + local = MicroSimulationLocal(3, True, MinimalSim) + # The outer local gid should remain 3 + self.assertEqual(local.get_global_id(), 3) + # The inner instance should have been constructed with -1 + self.assertEqual(local._instance.get_global_id(), -1) + + def test_initialize(self): + local = MicroSimulationLocal(0, False, SimWithInitialize) + result = local.initialize({"data": 1}) + self.assertEqual(result, {"init": True}) + + def test_output(self): + local = MicroSimulationLocal(0, False, SimWithOutput) + self.assertIsNone(local.output()) + + def test_requires_initialize(self): + local = MicroSimulationLocal(0, False, SimWithInitialize) + self.assertTrue(local.requires_initialize()) + + def test_requires_output(self): + local = MicroSimulationLocal(0, False, SimWithOutput) + self.assertTrue(local.requires_output()) + + def test_getattr_delegates_to_instance(self): + """__getattr__ should delegate unknown attributes to the wrapped instance.""" + + class SimWithExtra(MinimalSim): + extra_attr = "hello" + + local = MicroSimulationLocal(7, False, SimWithExtra) + # extra_attr is not defined on MicroSimulationLocal — must come via __getattr__ + self.assertEqual(local.extra_attr, "hello") + + +class TestMicroSimulationRemote(unittest.TestCase): + def _make_remote(self, late_init=False): + from micro_manager.micro_simulation import MicroSimulationRemote + + conn = MagicMock() + conn.recv.return_value = None + return ( + MicroSimulationRemote( + gid=0, + late_init=late_init, + num_ranks=1, + conn=conn, + cls_path="dummy_path", + sim_cls=MinimalSim, + ), + conn, + ) + + def test_get_set_global_id(self): + remote, _ = self._make_remote() + self.assertEqual(remote.get_global_id(), 0) + remote.set_global_id(42) + self.assertEqual(remote.get_global_id(), 42) + + def test_solve_returns_worker0_result(self): + remote, conn = self._make_remote() + conn.recv.return_value = {"out": 99} + result = remote.solve({"in": 1}, 0.1) + self.assertEqual(result, {"out": 99}) + + def test_get_state_returns_dict_keyed_by_worker(self): + remote, conn = self._make_remote() + conn.recv.return_value = {"gid": 0, "state": "s"} + state = remote.get_state() + self.assertIn(0, state) + self.assertEqual(state[0], {"gid": 0, "state": "s"}) + + def test_set_state_updates_gid_from_worker0(self): + remote, conn = self._make_remote() + conn.recv.return_value = 7 + remote.set_state({0: {"gid": 7, "state": "s"}}) + self.assertEqual(remote.get_global_id(), 7) + + def test_initialize_returns_worker0_result(self): + remote, conn = self._make_remote() + conn.recv.return_value = {"init": True} + result = remote.initialize() + self.assertEqual(result, {"init": True}) + + def test_output_returns_worker0_result(self): + remote, conn = self._make_remote() + conn.recv.return_value = None + result = remote.output() + self.assertIsNone(result) + + def test_requires_initialize_false_for_minimal_sim(self): + remote, _ = self._make_remote() + self.assertFalse(remote.requires_initialize()) + + def test_requires_initialize_true_for_sim_with_initialize(self): + from micro_manager.micro_simulation import MicroSimulationRemote + + conn = MagicMock() + conn.recv.return_value = None + remote = MicroSimulationRemote( + gid=0, + late_init=False, + num_ranks=1, + conn=conn, + cls_path="dummy_path", + sim_cls=SimWithInitialize, + ) + self.assertTrue(remote.requires_initialize()) + + def test_requires_output_false_for_minimal_sim(self): + remote, _ = self._make_remote() + self.assertFalse(remote.requires_output()) + + def test_requires_output_true_for_sim_with_output(self): + from micro_manager.micro_simulation import MicroSimulationRemote + + conn = MagicMock() + conn.recv.return_value = None + remote = MicroSimulationRemote( + gid=0, + late_init=False, + num_ranks=1, + conn=conn, + cls_path="dummy_path", + sim_cls=SimWithOutput, + ) + self.assertTrue(remote.requires_output()) + + def test_destroy_sends_delete_task_to_all_workers(self): + remote, conn = self._make_remote() + conn.recv.return_value = None + remote.destroy() + conn.send.assert_called() + conn.recv.assert_called() + + def test_late_init_uses_construct_late_task(self): + from micro_manager.micro_simulation import MicroSimulationRemote + from micro_manager.tasking.task import ConstructLateTask + + conn = MagicMock() + conn.recv.return_value = None + remote = MicroSimulationRemote( + gid=5, + late_init=True, + num_ranks=1, + conn=conn, + cls_path="dummy_path", + sim_cls=MinimalSim, + ) + sent_task = conn.send.call_args_list[0][0][1] + self.assertEqual(sent_task[0], "ConstructLateTask") + + +class TestCreateSimulationClass(unittest.TestCase): + def test_valid_class(self): + log = MagicMock() + sim_cls = create_simulation_class(log, MinimalSim, "dummy_path", 1) + self.assertIsNotNone(sim_cls) + + def test_missing_get_global_id_raises(self): + class BadSim: + def solve(self, i, dt): + pass + + def get_state(self): + pass + + def set_state(self, s): + pass + + log = MagicMock() + with self.assertRaises(ValueError): + create_simulation_class(log, BadSim, "dummy_path", 1) + + def test_missing_solve_raises(self): + class BadSim: + def get_global_id(self): + pass + + def get_state(self): + pass + + def set_state(self, s): + pass + + log = MagicMock() + with self.assertRaises(ValueError): + create_simulation_class(log, BadSim, "dummy_path", 1) + + def test_missing_get_state_raises(self): + class BadSim: + def get_global_id(self): + pass + + def set_state(self, s): + pass + + def solve(self, i, dt): + pass + + log = MagicMock() + with self.assertRaises(ValueError): + create_simulation_class(log, BadSim, "dummy_path", 1) + + def test_missing_set_state_raises(self): + class BadSim: + def get_global_id(self): + pass + + def get_state(self): + pass + + def solve(self, i, dt): + pass + + log = MagicMock() + with self.assertRaises(ValueError): + create_simulation_class(log, BadSim, "dummy_path", 1) + + def test_custom_sim_class_name(self): + log = MagicMock() + sim_cls = create_simulation_class( + log, MinimalSim, "dummy_path", 1, sim_class_name="MyTestSim" + ) + self.assertEqual(sim_cls.name, "MyTestSim") + + def test_interface_subclass_accepted_without_wrapping(self): + """A class that already inherits MicroSimulationInterface is accepted as-is.""" + log = MagicMock() + sim_cls = create_simulation_class(log, MinimalSim, "dummy_path", 1) + # backend_cls should be exactly MinimalSim — no wrapping applied + self.assertIs(sim_cls.backend_cls, MinimalSim) + + +class TestMicroSimulationClassMethods(unittest.TestCase): + def setUp(self): + self.log = MagicMock() + self.sim_cls = create_simulation_class(self.log, MinimalSim, "dummy_path", 1) + self.sim_cls_with_init = create_simulation_class( + self.log, SimWithInitialize, "dummy_path", 1 + ) + self.sim_cls_with_output = create_simulation_class( + self.log, SimWithOutput, "dummy_path", 1 + ) + + def test_check_output_false(self): + self.assertFalse(self.sim_cls.check_output()) + + def test_check_output_true(self): + self.assertTrue(self.sim_cls_with_output.check_output()) + + def test_check_initialize_false(self): + instance = MinimalSim(0) + has_init, has_args = self.sim_cls.check_initialize(instance, {}) + self.assertFalse(has_init) + self.assertFalse(has_args) + + def test_check_initialize_true_no_args(self): + class SimInitNoArgs(MinimalSim): + def initialize(self): + return None + + log = MagicMock() + sim_cls = create_simulation_class(log, SimInitNoArgs, "dummy", 1) + instance = SimInitNoArgs(0) + has_init, has_args = sim_cls.check_initialize(instance, {}) + self.assertTrue(has_init) + self.assertFalse(has_args) + + def test_check_initialize_true_with_args(self): + instance = SimWithInitialize(0) + has_init, has_args = self.sim_cls_with_init.check_initialize( + instance, {"data": 1} + ) + self.assertTrue(has_init) + self.assertTrue(has_args) + + def test_call_creates_wrapper(self): + wrapper = self.sim_cls(0) + self.assertIsNotNone(wrapper) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_micro_simulation_crash_handling.py b/tests/unit/test_micro_simulation_crash_handling.py index 6ddd7045..b6fbe1df 100644 --- a/tests/unit/test_micro_simulation_crash_handling.py +++ b/tests/unit/test_micro_simulation_crash_handling.py @@ -21,6 +21,15 @@ def solve(self, macro_data, dt): "Micro-Scalar-Data": macro_data["Macro-Scalar-Data"], } + def get_state(self): + return None + + def set_state(self, state): + pass + + def get_global_id(self): + return self.sim_id + class TestSimulationCrashHandling(TestCase): def test_crash_handling(self): diff --git a/tests/unit/test_snapshot_computation.py b/tests/unit/test_snapshot_computation.py index 73fec78d..55e7d617 100644 --- a/tests/unit/test_snapshot_computation.py +++ b/tests/unit/test_snapshot_computation.py @@ -19,6 +19,15 @@ def solve(self, macro_data, dt): "Micro-Vector-Data": macro_data["Macro-Vector-Data"] + 1, } + def get_state(self): + return None + + def set_state(self, state): + pass + + def get_global_id(self): + pass + class TestFunctionCalls(TestCase): def setUp(self): @@ -85,7 +94,11 @@ def test_solve_micro_sims(self): snapshot_object._micro_problem = MicroSimulation snapshot_object._micro_sims = create_simulation_class( - snapshot_object._micro_problem + MagicMock(), + snapshot_object._micro_problem, + None, + 1, + None, )(0) micro_sim_output = snapshot_object._solve_micro_simulation(self.fake_read_data) diff --git a/tests/unit/test_tasking.py b/tests/unit/test_tasking.py new file mode 100644 index 00000000..49a76409 --- /dev/null +++ b/tests/unit/test_tasking.py @@ -0,0 +1,208 @@ +from unittest import TestCase +from unittest.mock import MagicMock + +import os +import sys +from pathlib import Path + +# no clue why this TestCase does not see the local precice.py +# So we add path here and load it +sys.path.append(str(Path(__file__).resolve().parent)) +import precice + +import numpy as np +from mpi4py import MPI + +from micro_manager.micro_simulation import create_simulation_class +from micro_manager.tasking.connection import spawn_local_workers +from micro_manager.tasking.task import ( + ConstructTask, + ConstructLateTask, + InitializeTask, + OutputTask, + SolveTask, + SetStateTask, + GetStateTask, +) + +data_size = 32 +num_workers = 2 + + +class MicroSimulation: + def __init__(self, sim_id): + self._rank = MPI.COMM_WORLD.Get_rank() + self._state = "Some State" if sim_id >= 0 else "No State" + self._sim_id = sim_id + + def initialize(self): + return "initialized" + + def output(self): + return "output" + + def solve(self, macro_data, dt): + data = macro_data["task-data"] + + data_per_rank = data_size // num_workers + + data_local = data[data_per_rank * self._rank : data_per_rank * (self._rank + 1)] + sum_local = np.asarray(data_local).sum() + sum_global = MPI.COMM_WORLD.allreduce(sum_local, op=MPI.SUM) + result = { + "task-result": sum_global, + } + return result + + def get_global_id(self): + return self._sim_id + + def get_state(self): + return { + "gid": self._sim_id, + "state": self._state, + } + + def set_state(self, state): + assert "gid" in state + assert "state" in state + + self._sim_id = state["gid"] + self._state = state["state"] + + +class TestTasking(TestCase): + """ + Can only test for general functionality, not pinning. + """ + + def setUp(self): + # cannot use mpi, would need to set certain env flags + mm_dir = os.path.abspath(str(Path(__file__).resolve().parent.parent.parent)) + worker_exec = os.path.join(mm_dir, "micro_manager", "tasking", "worker_main.py") + + self.conn = spawn_local_workers( + worker_exec, num_workers, "socket", False, "open", "" + ) + self.input_data = {"task-data": np.arange(data_size)} + self.expected_output = (data_size - 1) * data_size / 2 + self.cls_path = os.path.abspath(str(Path(__file__).resolve())).replace( + ".py", "" + ) + self.sim_cls = create_simulation_class( + MagicMock(), + MicroSimulation, + self.cls_path, + num_workers, + self.conn, + ) + + def tearDown(self): + super().tearDown() + if self.conn is not None: + self.conn.close() + + def test_construct(self): + gid = 0 + + self.send(ConstructTask.send_args(gid, self.cls_path)) + self.recv() + + def test_construct_late(self): + gid = 0 + + self.send(ConstructLateTask.send_args(gid, self.cls_path)) + self.recv() + + base_state = { + "gid": gid, + "state": "Important State Information", + } + states = [base_state, base_state] + for i in range(num_workers): + self.conn.send(i, SetStateTask.send_args(gid, states[i])) + self.recv() + + def test_initialize(self): + gid = 0 + sim = self.sim_cls(gid) + + self.send(InitializeTask.send_args(gid)) + result = self.recv() + for i in range(num_workers): + self.assertEqual(result[i], "initialized") + + def test_output(self): + gid = 0 + sim = self.sim_cls(gid) + + self.send(OutputTask.send_args(gid)) + result = self.recv() + for i in range(num_workers): + self.assertEqual(result[i], "output") + + def test_solve(self): + gid = 0 + sim = self.sim_cls(gid) + + result_interface = sim.solve(self.input_data, 0.0) + self.assertEqual(result_interface["task-result"], self.expected_output) + + self.send(SolveTask.send_args(gid, self.input_data, 0.0)) + result_manual = self.recv() + self.assertTrue(result_manual[0] is not None) + self.assertDictEqual(result_interface, result_manual[0]) + + def test_get_state(self): + gid = 0 + sim = self.sim_cls(gid) + + result_interface = list(sim.get_state().values()) + for i in range(num_workers): + self.assertEqual(result_interface[i]["gid"], gid) + self.assertEqual(result_interface[i]["state"], "Some State") + + self.send(GetStateTask.send_args(gid)) + result_manual = self.recv() + self.assertListEqual(result_manual, result_interface) + + def test_set_state(self): + gid = -1 + gid_new = 0 + sim = self.sim_cls(gid) + base_state = { + "gid": gid_new, + "state": "Important State Information", + } + states = [base_state, base_state] + + sim.set_state(states) + result_interface = sim.get_state() + self.assertEqual(result_interface[0]["gid"], gid_new) + self.assertEqual(result_interface[0]["state"], "Important State Information") + + sim.destroy() + del sim + + sim = self.sim_cls(gid) + for i in range(num_workers): + self.conn.send(i, SetStateTask.send_args(gid, states[i])) + self.recv() + + self.send(GetStateTask.send_args(gid_new)) + result_manual = self.recv() + self.assertEqual(result_manual[0]["gid"], gid_new) + self.assertEqual(result_manual[0]["state"], "Important State Information") + + def send(self, args): + for i in range(num_workers): + self.conn.send(i, args) + + def recv(self): + return [self.conn.recv(i) for i in range(num_workers)] + + +if __name__ == "__main__": + import unittest + + unittest.main()