Enable colony simulation state saving and custom emit_step configuration - #442
Open
katha815 wants to merge 6 commits into
Open
Enable colony simulation state saving and custom emit_step configuration#442katha815 wants to merge 6 commits into
katha815 wants to merge 6 commits into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR enables proper colony simulation state saving and Parquet emission finalization for both dividing and non-dividing runs. It also ensures the
emit_stepconfiguration is correctly propagated to both outer and inner emitters.All tests were run with the following command structure:
Fix Details
Fix 0 — Convert Two-Component System to Vivarium
PartitionedProcessWhile investigating the colony state saving issue (Fix 1), I found that the Two-Component System process was still using the standalone interface and was not fully compatible with the current Vivarium framework. This may cause compatibility issues when the simulation attempted to save and restore colony states.
To address this, I migrated the process to
PartitionedProcess:topology_registry.ports_schema(),calculate_request(), andevolve_state()methods.bulk_name_to_idx()andcounts()helpers.This migration was necessary for the colony simulation to handle states correctly through the
PartitionedProcessinterface.Fix 1 — Fix colony final-state JSON saving
configs/colony_baseline_test.json. The colony's final state JSON was not being saved correctly. The Parquet emitter wraps agents in anagents/outerstructure, butcolony_save_states()expectedagents/<agent_id>/...directly.agents/outer/agents/<agent_id>/cell_processagents/<agent_id>/cell_processstate_to_save["agents"]["outer"]before saving so colony JSON retains expected structure.bulk_dtypesandunique_dtypesmetadata.configs/colony_baseline_test.json(2-generation dividing run, 1-second smoke test)agents/<agent_id>/...structure.Fix 2 — Finalize outer Parquet emitter on shutdown
configs/colony_baseline_test3.json. The emission step remained at 1 second regardless of custom configuration. The outerParquetEmitterbuffers writes in batches of 400 emits, and runs that don't reach this threshold leave buffered data unwritten.engine.end()inrun_simulation(), mark the outer Parquet emitter as successful and explicitly callengine.emitter.finalize().configs/colony_baseline_test4.json(dividing run, 60-second emit step)Fix 3 — Finalize inner Parquet emitters and propagate
emit_step/out/{project}/...) saved Parquet data, but the inner subfolder (/out/{project}/...__inner) either saved data every 1 second or did not save at all. This occurred because the innerParquetEmitterwas only finalized from division or exception paths insideEngineProcess.next_update(), so short non-dividing runs left inner Parquet data unwritten.emit_stepfrom top-level simulation configuration throughEcoliEngineProcessinto the innerEngineProcess.finalize_parquet_emitters()helper that recursively traverses nested process dictionaries and callsfinalize()on everyParquetEmitterwhile avoiding duplicate finalization.configs/colony_baseline_test4.json(short non-dividing run, 60-second emit step)emit_steppropagation and proper finalization for non-dividing runs.Test Configurations Used
configs/colony_baseline_test.jsonconfigs/colony_baseline_test3.jsonconfigs/colony_baseline_test4.jsonAll three test configurations produced the expected outputs. The final JSON from Fix 1 has the correct
agents/<agent_id>/...structure. The Parquet outputs from Fix 2 and Fix 3 contain emissions at 0s, 60s, and 120s.Files Changed
ecoli/experiments/ecoli_engine_process.py— Main fixes for state saving and emitter finalizationPartitionedProcessconfigs/colony_baseline_test.json,configs/colony_baseline_test3.json,configs/colony_baseline_test4.json— Test configurationsNotes
The current implementation prioritises correctness over code brevity. Additional simplification may be considered in a follow-up.