pyalma/al2/plan.py

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()