Skip to content
Snippets Groups Projects

Draft: Support restoring the world state by setting the RESTORE_PATH environment variable.

Open Adam Reichold requested to merge restore-state into main
Files
17
+ 71
4
#![feature(write_all_vectored)]
use std::fs::File;
use std::env::var_os;
use std::fs::{read_dir, File};
use std::io::{IoSlice, Write};
use std::mem::transmute;
use std::ops::Deref;
use std::path::Path;
use memmap2::MmapOptions;
use rand_distr::Bernoulli;
use dump::save;
use model::{params::Params, units::Tick, Model};
use dump::{restore, save};
use model::{params::Params, units::Tick, Model, World};
use process::run;
pub fn preferences<'a, P: IntoIterator<Item = &'a f64>>(probs: P) -> Box<[Bernoulli]> {
@@ -23,6 +25,12 @@ pub fn sample(params: Params, duration: Tick, interval: Tick) {
run(move |rank, size, process| {
let mut model = Model::new(params, rank, size);
let first_snapshot = if let Some(restore_path) = var_os("RESTORE_PATH") {
restore_world(restore_path, rank, &mut model.world)
} else {
0
};
let mut buffer = Vec::new();
let locations = model.flower_patch_locations.deref();
@@ -49,11 +57,22 @@ pub fn sample(params: Params, duration: Tick, interval: Tick) {
let duration = duration / interval;
let interval = interval / model.params.tick;
for snapshot in 0..duration {
for snapshot in first_snapshot..duration {
for _ in 0..interval {
model.run(process);
}
if model.params.trace_temporal_network {
let tick = model.world.tick;
for flower_patch in &mut *model.world.flower_patches {
for flower_contaminators in &mut *flower_patch.flower_contaminators {
flower_contaminators
.retain(|(_contaminator, decontaminates_at)| *decontaminates_at > tick);
}
}
}
write_vectored(
format!("world_{:03}_{:02}.bin", snapshot, model.grid.rank),
&mut buffer,
@@ -67,6 +86,54 @@ pub fn sample(params: Params, duration: Tick, interval: Tick) {
});
}
fn restore_world<P>(path: P, rank: u8, world: &mut World) -> i64
where
P: AsRef<Path>,
{
let prefix = "world_";
let suffix = format!("_{rank:02}.bin");
let mut last_snapshot = None;
let mut last_snapshot_path = None;
for entry in read_dir(path).unwrap() {
let entry = entry.unwrap();
if let Some(file_name) = entry.file_name().to_str() {
if file_name.starts_with(prefix) && file_name.ends_with(&suffix) {
let snapshot = file_name.split('_').nth(1).unwrap().parse::<i64>().unwrap();
match &mut last_snapshot {
None => {
last_snapshot = Some(snapshot);
last_snapshot_path = Some(entry.path());
}
Some(last_snapshot) if *last_snapshot < snapshot => {
*last_snapshot = snapshot;
last_snapshot_path = Some(entry.path());
}
Some(_) => (),
}
}
}
}
let last_snapshot = last_snapshot.unwrap();
let last_snapshot_path = last_snapshot_path.unwrap();
eprintln!("Restoring snapshot {last_snapshot} of rank {rank}...");
unsafe {
let mut buf = MmapOptions::new()
.map_copy(&File::open(last_snapshot_path).unwrap())
.unwrap();
*world = restore::<World>(&mut buf).clone();
}
last_snapshot + 1
}
fn write_vectored<'a, P, F>(path: P, buffer: &'a mut Vec<IoSlice<'static>>, f: F)
where
P: AsRef<Path>,
Loading