224 lines
7.7 KiB
Python
224 lines
7.7 KiB
Python
#! /bin/python3
|
|
|
|
import time, sys, random, copy
|
|
import importlib
|
|
import subprocess
|
|
from joblib import Parallel, delayed
|
|
from DAS import *
|
|
import os
|
|
import pickle
|
|
import uuid
|
|
|
|
# Parallel execution:
|
|
# The code currently uses 'joblib' to execute on multiple cores. For other options such as 'ray', see
|
|
# https://stackoverflow.com/questions/9786102/how-do-i-parallelize-a-simple-python-loop
|
|
# For fixing logging issues in parallel execution, see
|
|
# https://stackoverflow.com/questions/58026381/logging-nested-functions-using-joblib-parallel-and-delayed-calls
|
|
# and https://github.com/joblib/joblib/issues/1017
|
|
|
|
def initLogger(config):
|
|
"""It initializes the logger."""
|
|
logger = logging.getLogger("Study")
|
|
logger.setLevel(config.logLevel)
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(config.logLevel)
|
|
ch.setFormatter(CustomFormatter())
|
|
logger.addHandler(ch)
|
|
return logger
|
|
|
|
def runOnce(config, shape, execID):
|
|
|
|
if config.deterministic:
|
|
shape.setSeed(config.randomSeed+"-"+str(shape))
|
|
random.seed(shape.randomSeed)
|
|
|
|
backup_folder = f"results/{execID}/backup"
|
|
if not os.path.exists(backup_folder):
|
|
os.makedirs(backup_folder)
|
|
backup_file = os.path.join(backup_folder, f"simulation_data_{shape}.pkl")
|
|
with open(backup_file, 'ab') as f:
|
|
pickle.dump(shape.__dict__, f)
|
|
|
|
sim = Simulator(shape, config, execID)
|
|
sim.initLogger()
|
|
sim.initValidators()
|
|
sim.initNetwork()
|
|
result = sim.run()
|
|
sim.logger.info("Shape: %s ... Block Available: %d in %d steps" % (str(sim.shape.__dict__), result.blockAvailable, len(result.missingVector)), extra=sim.format)
|
|
|
|
if config.dumpXML:
|
|
result.dump()
|
|
|
|
if config.visualization:
|
|
visual = Visualizor(execID, config, [result])
|
|
visual.plotAll()
|
|
|
|
with open(backup_file, 'ab') as f:
|
|
pickle.dump("completed", f)
|
|
|
|
return result
|
|
|
|
|
|
def check_simulation_completion(state_file):
|
|
backup_dir = os.path.join(os.path.dirname(state_file), "backup")
|
|
if not os.path.exists(backup_dir):
|
|
return False
|
|
|
|
all_completed = True
|
|
incomplete_files = []
|
|
completed_files = []
|
|
completed_shapes = []
|
|
for filename in sorted(os.listdir(backup_dir), reverse=True):
|
|
if not filename.endswith(".pkl"):
|
|
continue
|
|
full_path = os.path.join(backup_dir, filename)
|
|
try:
|
|
with open(full_path, 'rb') as f:
|
|
items = []
|
|
while True:
|
|
try:
|
|
item = pickle.load(f)
|
|
items.append(item)
|
|
except EOFError:
|
|
break
|
|
last_item = items[-1]
|
|
if last_item != "completed":
|
|
all_completed = False
|
|
incomplete_files.append(full_path)
|
|
else:
|
|
completed_files.append(full_path)
|
|
completed_shapes.append(items[0])
|
|
except (OSError, pickle.UnpicklingError) as e:
|
|
print(f"Error loading state from {full_path}: {e}")
|
|
all_completed = False
|
|
break
|
|
return all_completed, incomplete_files, completed_files, completed_shapes
|
|
|
|
|
|
def start_simulation(execID, completed_files, completed_shapes, incomplete_files):
|
|
config = importlib.import_module("smallConf")
|
|
logger = initLogger(config)
|
|
format = {"entity": "Study"}
|
|
|
|
results = []
|
|
if not os.path.exists("results"):
|
|
os.makedirs("results")
|
|
dir = "results/"+execID
|
|
if not os.path.exists(dir):
|
|
os.makedirs(dir)
|
|
if config.saveGit:
|
|
with open(dir+"/git.diff", 'w') as f:
|
|
subprocess.run(["git", "diff"], stdout=f)
|
|
with open(dir+"/git.describe", 'w') as f:
|
|
subprocess.run(["git", "describe", "--always"], stdout=f)
|
|
|
|
logger.info("Starting simulations:", extra=format)
|
|
start = time.time()
|
|
for shape in config.nextShape():
|
|
comparison_dict = shape.__dict__.copy()
|
|
ignore_keys = ['randomSeed']
|
|
for key in ignore_keys:
|
|
del comparison_dict[key]
|
|
|
|
if any(all(comparison_dict[key] == completed_shape[key] for key in comparison_dict.keys() if key not in ignore_keys) for completed_shape in completed_shapes):
|
|
logger.info("Skipping simulation for shape (already completed): %s" % (str(shape.__dict__)), extra=format)
|
|
else:
|
|
results.append(delayed(runOnce)(config, shape, execID))
|
|
|
|
results = Parallel(config.numJobs)(results)
|
|
end = time.time()
|
|
logger.info("A total of %d simulations ran in %d seconds" % (len(results), end-start), extra=format)
|
|
|
|
if config.visualization:
|
|
vis = Visualizer(execID, config)
|
|
vis.plotHeatmaps()
|
|
|
|
visual = Visualizor(execID, config, results)
|
|
visual.plotHeatmaps("nn", "fr")
|
|
|
|
|
|
def study():
|
|
restart_path = None
|
|
for arg in sys.argv[1:]:
|
|
if arg.startswith("--restart="):
|
|
restart_path = arg[len("--restart="):]
|
|
|
|
if restart_path:
|
|
execID = restart_path.split("/")[1]
|
|
state_file = f"results/{execID}/backup"
|
|
all_completed, incomplete_files, completed_files, completed_shapes = check_simulation_completion(state_file)
|
|
|
|
current_shapes = []
|
|
config = importlib.import_module("smallConf")
|
|
|
|
completed_shapes_without_seed = completed_shapes
|
|
for shape in config.nextShape():
|
|
shape_dict = copy.deepcopy(shape.__dict__)
|
|
del shape_dict['randomSeed']
|
|
current_shapes.append(shape_dict)
|
|
for shape in completed_shapes_without_seed:
|
|
if 'randomSeed' in shape:
|
|
del shape['randomSeed']
|
|
|
|
completed_set = {frozenset(shape.items()) for shape in completed_shapes_without_seed}
|
|
current_set = {frozenset(shape.items()) for shape in current_shapes}
|
|
|
|
if all_completed and completed_set == current_set:
|
|
print("Simulation is already completed.")
|
|
sys.exit(0)
|
|
else:
|
|
print("Restarting simulations.")
|
|
start_simulation(execID, completed_files, completed_shapes, incomplete_files)
|
|
sys.exit(0)
|
|
|
|
if len(sys.argv) < 2:
|
|
print("You need to pass a configuration file in parameter")
|
|
exit(1)
|
|
|
|
try:
|
|
config = importlib.import_module(sys.argv[1])
|
|
except ModuleNotFoundError as e:
|
|
try:
|
|
config = importlib.import_module(str(sys.argv[1]).replace(".py", ""))
|
|
except ModuleNotFoundError as e:
|
|
print(e)
|
|
print("You need to pass a configuration file in parameter")
|
|
exit(1)
|
|
|
|
logger = initLogger(config)
|
|
format = {"entity": "Study"}
|
|
|
|
results = []
|
|
|
|
now = datetime.now()
|
|
execID = now.strftime("%Y-%m-%d_%H-%M-%S_")+str(random.randint(100,999))
|
|
|
|
# save config and code state for reproducibility
|
|
if not os.path.exists("results"):
|
|
os.makedirs("results")
|
|
dir = "results/"+execID
|
|
if not os.path.exists(dir):
|
|
os.makedirs(dir)
|
|
if config.saveGit:
|
|
with open(dir+"/git.diff", 'w') as f:
|
|
subprocess.run(["git", "diff"], stdout=f)
|
|
with open(dir+"/git.describe", 'w') as f:
|
|
subprocess.run(["git", "describe", "--always"], stdout=f)
|
|
subprocess.run(["cp", sys.argv[1], dir+"/"])
|
|
|
|
logger.info("Starting simulations:", extra=format)
|
|
start = time.time()
|
|
results = Parallel(config.numJobs)(delayed(runOnce)(config, shape ,execID) for shape in config.nextShape())
|
|
end = time.time()
|
|
logger.info("A total of %d simulations ran in %d seconds" % (len(results), end-start), extra=format)
|
|
|
|
if config.visualization:
|
|
vis = Visualizer(execID, config)
|
|
vis.plotHeatmaps()
|
|
|
|
visual = Visualizor(execID, config, results)
|
|
visual.plotHeatmaps("nn", "fr")
|
|
|
|
if __name__ == "__main__":
|
|
study()
|