import pathlib as pl import json import os import multiprocessing import threading from . import batch class Plan: def __init__(self, experiment=None, lock=None): self.experiment = None self.file = None self.pending_instances = [] self.assigned_instances = [] self.__lock = threading.Lock() if lock == None else lock if experiment: self.create(experiment) def create(self, experiment): self.experiment = pl.Path(experiment).resolve() self.__set_file() if self.__is_finished(): self.__create() else: self.__load() def __create(self): content = self.__create_content() self.pending_instances = content["pending"] self.iterations_left = content["iterations_left"] with self.__lock: self.__update_file() def __create_content(self, iterations_left = None): content = {} with open(self.experiment, "r") as expf: exp_obj = json.loads(expf.read()) instances = batch.load(pl.Path(exp_obj["batch"]).resolve()) if iterations_left == None: if "iterations" in exp_obj: iterations_left = exp_obj["iterations"] - 1 else: iterations_left = 0 content["pending"] = instances content["iterations_left"] = iterations_left return content def __set_file(self): if self.experiment == None: self.file = None else: exp_path = pl.Path(self.experiment) self.file = exp_path.parent / (exp_path.stem + ".plan") def __load(self): self.pending_instances = [] self.assigned_instances = [] if not self.file.is_file(): return with open(self.file, "r") as pfile: content = json.loads(pfile.read()) if "assigned" in content: self.assigned_instances = content["assigned"] if "pending" in content: self.pending_instances = content["pending"] if "iterations_left" in content: self.iterations_left = content["iterations_left"] def __is_finished(self): return False if self.file.is_file() else True def next(self): with self.__lock: self.__load() if len(self.pending_instances) == 0: if self.iterations_left > 0: self.__load_next_iteration() else: return None next_instance = self.pending_instances.pop() self.assigned_instances.append(next_instance) self.__update_file() return next_instance def done_with(self, instance): with self.__lock: self.__load() if instance in self.assigned_instances: self.assigned_instances.remove(instance) self.__update_file() def __update_file(self): content = {} all_done = True content["iterations_left"] = self.iterations_left if len(self.assigned_instances) > 0: content["assigned"] = self.assigned_instances all_done = False if len(self.pending_instances) > 0: content["pending"] = self.pending_instances all_done = False if all_done: if self.iterations_left > 0: self.__load_next_iteration() elif self.file.is_file(): self.file.unlink() else: self.__write_content(content) def __load_next_iteration(self): content = self.__create_content(self.iterations_left - 1) self.pending_instances = content["pending"] self.iterations_left = content["iterations_left"] self.__write_content(content) def __write_content(self, content): if "assigned" in content: content["assigned"][:] = map(str, content["assigned"]) if "pending" in content: content["pending"][:] = map(str, content["pending"]) with open(self.file, "w") as pfile: pfile.write(json.dumps(content)) def delete(self): with self.__lock: self.__load() self.pending_instances.extend(self.assigned_instances) self.assigned_instances = [] self.__update_file()