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): with open(self.experiment, "r") as expf: exp_obj = json.loads(expf.read()) instances = batch.load(pl.Path(exp_obj["batch"]).resolve()) self.pending_instances = instances self.__update_file() 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"] 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: 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 = {} if len(self.assigned_instances) > 0: content["assigned"] = list(map(str, self.assigned_instances)) if len(self.pending_instances) > 0: content["pending"] = list(map(str, self.pending_instances)) if content: with open(self.file, "w") as pfile: pfile.write(json.dumps(content)) elif self.file.is_file(): self.file.unlink() #def __del__(self): def delete(self): with self.__lock: self.__load() self.pending_instances.extend(self.assigned_instances) self.assigned_instances = [] self.__update_file()