diff --git a/doc/docs/Parallel_Meep.md b/doc/docs/Parallel_Meep.md index ba77dcfc7..3a4452179 100644 --- a/doc/docs/Parallel_Meep.md +++ b/doc/docs/Parallel_Meep.md @@ -189,7 +189,7 @@ The chunk balancer interface provides four main methods: - `adjust_chunk_layout()` is syntactic sugar which will compute a new chunk layout, apply it to the simulation object, reset the simulation, and re-initialize the simulation. ```py -class ChunkyMonkey(abc.ABC): +class AbstractChunkBalancer(abc.ABC): """Chunk balancer for dynamically load-balanced chunk layouts.""" def make_initial_chunk_layout(self, sim) -> mp.BinaryPartition: @@ -216,9 +216,9 @@ class ChunkyMonkey(abc.ABC): Using the chunk balancer is very straightforward, and it can typically be integrated into existing Meep simulations with only a few lines of code. Here is a simple example: ```py -from meep.chunk_balancer import DefaultChunkBalancer +from meep.chunk_balancer import ChunkBalancer -chunk_balancer = DefaultChunkBalancer() +chunk_balancer = ChunkBalancer() # Compute an initial chunk layout initial_chunk_layout = chunk_balancer.make_initial_chunk_layout() @@ -228,11 +228,42 @@ sim.init_sim() for iteration in range(epochs): sim.run(...) - - # Adjust the chunk layout if needed + # Adjust the chunk layout for the next iteration if needed chunk_balancer.adjust_chunk_layout(sim, sensitivity=0.4) ``` +Chunks can also be rebalanced between runs of a program by dumping and loading the chunk layout from a pickled object. For example: + +```py +import meep as mp +from meep.chunk_balancer import ChunkBalancer +from meep.timing_measurements import MeepTimingMeasurements +import pickle + +# Fetch chunk layout from a previous run +with open("path/to/chunk_layout.pkl", "rb") as f: + chunk_layout = pickle.load(f) + +sim = mp.Simulation(..., chunk_layout=chunk_layout) +sim.init_sim() +sim.run(...) + +# Compute and save chunk layout for next run +timings = MeepTimingMeasurements.new_from_simulation(sim) +chunk_volumes = sim.structure.get_chunk_volumes() +chunk_owners = sim.structure.get_chunk_owners() +next_chunk_layout = ChunkBalancer().compute_new_chunk_layout( + timings, + chunk_layout, + chunk_volumes, + chunk_owners, + sensitivity=0.4) + +# Save chunk layout for next run +with open("path/to/chunk_layout.pkl", "wb") as f: + pickle.dump(next_chunk_layout, f) +``` + ### Chunk adjustment algorithm The default chunk adjustment algorithm recursively traverses the `BinaryPartition` object and resizes the chunk volumes of the left and right children of each node to have an equal per-node simulation time. The new chunk layout has split positions which are a weighted average of the previous iteration's chunk layout and the newly computed layout, and the `sensitivity` parameter adjusts how fast the chunk sizes are adjusted. (`sensitivity=0.0` means no adjustment, `sensitivity=1.0` means an immediate snap to the predicted split positions, and `sensitivity=0.5` is the average of the old and new layouts.) The adjustment algorithm is summarized in pseudocode below: diff --git a/python/chunk_balancer.py b/python/chunk_balancer.py index b109e480f..956ba979c 100644 --- a/python/chunk_balancer.py +++ b/python/chunk_balancer.py @@ -9,7 +9,7 @@ import numpy as np -class ChunkyMonkey(abc.ABC): +class AbstractChunkBalancer(abc.ABC): """Abstract chunk balancer for adaptive chunk layouts in Meep simulations. This class defines interfaces for a chunk balancer, which adjusts chunk @@ -135,7 +135,7 @@ def _validate_sim(self, sim: mp.Simulation): set(proc_ids) - set(chunk_owners))) -class DefaultChunkBalancer(ChunkyMonkey): +class ChunkBalancer(AbstractChunkBalancer): """A chunk balancer for adaptive chunk layouts in Meep simulations. This class generates initial chunk layouts using Meep's built-in scheme, and diff --git a/python/tests/test_chunk_balancer.py b/python/tests/test_chunk_balancer.py index 547422c27..36a85f40c 100644 --- a/python/tests/test_chunk_balancer.py +++ b/python/tests/test_chunk_balancer.py @@ -7,7 +7,7 @@ import meep.binary_partition_utils as bpu from meep.timing_measurements import MeepTimingMeasurements, TIMING_MEASUREMENT_IDS -from meep.chunk_balancer import DefaultChunkBalancer +from meep.chunk_balancer import ChunkBalancer class MockSimulation(mp.Simulation): @@ -200,7 +200,7 @@ def test_structure_get_chunk_owners(self, test_sim_constructor, self.assertListEqual(list(chunk_owners), expected_chunk_owners) -class DefaultChunkBalancerTest(unittest.TestCase): +class ChunkBalancerTest(unittest.TestCase): @parameterized.parameterized.expand([ (TEST_SIM_1, False), @@ -210,7 +210,7 @@ def test_validate_sim(self, test_sim_constructor, should_raise_exception): test_sim = test_sim_constructor() test_sim.init_sim() - chunk_balancer = DefaultChunkBalancer() + chunk_balancer = ChunkBalancer() if should_raise_exception: with self.assertRaises(ValueError): @@ -232,7 +232,7 @@ def test_chunk_layout_improvement(self, test_sim_constructor): old_timing_measurements = MeepTimingMeasurements.new_from_simulation( test_sim, -1) - chunk_balancer = DefaultChunkBalancer() + chunk_balancer = ChunkBalancer() chunk_balancer.adjust_chunk_layout(test_sim, sensitivity=1.0) @@ -264,7 +264,7 @@ def test_chunk_layout_convergence(self, test_sim_constructor): test_sim = test_sim_constructor() test_sim.init_sim() - chunk_balancer = DefaultChunkBalancer() + chunk_balancer = ChunkBalancer() num_iterations = 25 @@ -301,7 +301,7 @@ def test_split_pos_adjustment(self, test_chunk_data): chunk_volumes = sim.structure.get_chunk_volumes() chunk_owners = sim.structure.get_chunk_owners() - chunk_balancer = DefaultChunkBalancer() + chunk_balancer = ChunkBalancer() measurements = {} num_procs = len(test_chunk_data["time_stepping"]) diff --git a/python/tests/test_timing_measurements.py b/python/tests/test_timing_measurements.py index 768824193..e84e76105 100644 --- a/python/tests/test_timing_measurements.py +++ b/python/tests/test_timing_measurements.py @@ -15,11 +15,7 @@ def test_timing_measurements(self): ) time_start = time.time() sim.run(until=5) - timing_measurements = timing.MeepTimingMeasurements.new_from_simulation( - sim=sim, - elapsed_time=time.time() - time_start, - time_per_step=[0.1], - ) + timing_measurements = timing.MeepTimingMeasurements.new_from_simulation(sim) # Check for expected names after updating self.assertSetEqual( diff --git a/python/timing_measurements.py b/python/timing_measurements.py index 04aada9ef..08fc068b6 100644 --- a/python/timing_measurements.py +++ b/python/timing_measurements.py @@ -93,7 +93,7 @@ def __init__(self, def new_from_simulation( cls, sim: mp.Simulation, - elapsed_time: float, + elapsed_time: Optional[float] = -1, time_per_step: Optional[List[float]] = None, dft_relative_change: Optional[List[float]] = None, overlap_relative_change: Optional[List[float]] = None,