Use Process Commands for Saving States

Requires PR vivarium-collective/vivarium-core#198, which introduces
process commands. The process commands interface lets us retrieve
internal state from parallelized EngineProcess instances.

ecoli/composites/ecoli_engine_process.py

@@ -123,11 +123,15 @@ def not_a_process(value):
         del state_to_save['agents']  # Replace 'agents' with agent states
         state_to_save['agents'] = {}
 
+        # Get internal state from the EngineProcess sub-simulation
+        for agent_id in state['agents']:
+            engine.state.get_path(
+                ('agents', agent_id, 'cell_process')
+            ).value.send_command('get_inner_state')
         for agent_id in state['agents']:
-            # Get internal state from the EngineProcess sub-simulation
             cell_state = engine.state.get_path(
                 ('agents', agent_id, 'cell_process')
-            ).value.sim.state.get_value(condition=not_a_process)
+            ).value.get_command_result()
             del cell_state['environment']['exchange_data']  # Can't save, but will be restored when loading state
             state_to_save['agents'][agent_id] = cell_state
 

ecoli/processes/engine_process.py

@@ -65,7 +65,7 @@
 from vivarium.core.engine import Engine
 from vivarium.core.process import Process, Step
 from vivarium.core.registry import updater_registry, divider_registry
-from vivarium.core.store import DEFAULT_SCHEMA
+from vivarium.core.store import DEFAULT_SCHEMA, Store
 from vivarium.library.topology import get_in
 
 from ecoli.library.sim_data import RAND_MAX
@@ -233,6 +233,27 @@ def calculate_timestep(self, states):
             timestep = min(timestep, process.calculate_timestep({}))
         return timestep
 
+    def get_inner_state(self) -> None:
+        def not_a_process(value):
+            return not (isinstance(value, Store) and value.topology)
+        return self.sim.state.get_value(condition=not_a_process)
+
+    def send_command(self, command, args = None, kwargs = None) -> None:
+        if self._pending_command:
+            raise RuntimeError(
+                f'Trying to send command {(command, args, kwargs)} but '
+                f'command {self._pending_command} is still pending.')
+        self._pending_command = command, args, kwargs
+
+        args = args or tuple()
+        kwargs = kwargs or {}
+
+        if command == 'get_inner_state':
+            self._command_result = self.get_inner_state()
+        else:
+            self._pending_command = None
+            super().send_command(command, args, kwargs)
+
     def next_update(self, timestep, states):
         # Check whether we are being forced to finish early. This check
         # should happen before we mutate the inner simulation state to

requirements.txt

@@ -128,7 +128,7 @@ typing_extensions==4.0.1
 Unum==4.1.4
 vivarium-bioscrape==0.0.0.7
 vivarium-convenience==0.0.3
-vivarium-core==1.2.8
+vivarium-core==1.3.0
 vivarium-multibody==0.0.16
 wcwidth==0.2.5
 webencodings==0.5.1