class Experiment(object):

    @classmethod
    def name_computation(cls, exp_name, index):
        return "Computation-{}-{:d}".format(exp_name, index)

    def __init__(self, exp_name, parameter_set, computation_factory,
                 storage_factory=PickleStorage, user=None):
        self.exp_name = exp_name
        self.parameter_set = parameter_set
        self.comp_factory = computation_factory
        self.storage_factory = storage_factory
        self.monitor = Monitor(exp_name, user, storage_factory)

    @property
    def storage(self):
        return self.monitor.storage

    def yield_computations(self, context="n/a", start=0, capacity=None,
                           auto_refresh=False):
        if capacity is None:
            capacity = sys.maxsize

        self.monitor.refresh()
        storage = self.monitor.storage
        storage.init()
        self.storage.save_parameter_set(self.parameter_set)
        unlaunchable = self.monitor.unlaunchable_comp_names()

        storage_factory = self.storage_factory

        i = 0
        for j, param_dict in self.parameter_set:
            if j < start:
                continue
            if i >= capacity:
                break
            if auto_refresh:
                self.monitor.refresh()
                unlaunchable = self.monitor.unlaunchable_comp_names()

            label = Experiment.name_computation(self.exp_name, j)
            if label in unlaunchable:
                continue

            computation = self.comp_factory(exp_name=self.exp_name,
                                            comp_name=label,
                                            context=context,
                                            storage_factory=storage_factory)

            yield computation.lazyfy(**param_dict)

            i += 1

    def __iter__(self):
        for x in self.yield_computations():
            yield x

    def __len__(self):
        return len(self.parameter_set) class ParameterSet(CartesianParameterSet):
    pass class InSituEnvironment(Environment):
    """
    InSituEnvironment
    =================
    An `Environment` that runs the code in a sequential way, in the current
    process.

    Constructor Parameters
    ----------------------
    stdout: boolean (defaul: False)
        Whether to print on the standard output rather than redirect
        to a log file. Setting False is the standard way for environment.
    fail_fast: boolean (default: True
    """
    @classmethod
    def is_usable(cls):
        return True

    @property
    def auto_refresh(self):
        return True

    def __init__(self, stdout=False, fail_fast=True):
        super().__init__(fail_fast)
        self.stdout = stdout

    def __repr__(self):
        return "{cls}(stdout={stdout}, fail_fast={fail_fast})"\
               "".format(cls=self.__class__.__name__,
                         stdout=repr(self.stdout),
                         fail_fast=self.fail_fast)

    def issue(self, lazy_computation):
        if self.stdout:
            lazy_computation()
            return

        storage = lazy_computation.storage
        log_file = storage.get_log_prefix(lazy_computation.comp_name,
                                          ".{}".format(str(epoch())))
        sys_backup = sys.stdout, sys.stderr
        try:
            with open(log_file, "w") as hdl:
                sys.stdout = sys.stderr = hdl
                lazy_computation()
        finally:
            sys.stdout, sys.stderr = sys_backup class BashEnvironment(Environment):
    # For debugging purpose

    @classmethod
    def is_usable(cls):
        return True

    def __init__(self, serializer=Serializer(), fail_fast=True):
        super(BashEnvironment, self).__init__(fail_fast)
        # Since Serializer is totally stateless, it can be shared among
        # several instances
        self.serializer = serializer

    def __repr__(self):
        return "{cls}(serializer={serializer}, fail_fast={fail_fast})" \
               "".format(cls=self.__class__.__name__,
                         fail_fast=repr(self.fail_fast),
                         serializer=repr(self.serializer))

    def issue(self, lazy_computation):
        job_command = self.serializer.serialize_and_script(lazy_computation)

        storage = lazy_computation.storage
        log_file = storage.get_log_prefix(lazy_computation.comp_name,
                                          "-{}".format(str(epoch())))

        # Dirty hack: since we don't want to wait for the computation to end
        # for closing the file, we leave it as is.
        # For sequential computation `InSituEnvironment` is much cleaner
        file_handle = open(log_file, "w")
        subprocess.Popen(job_command, stderr=file_handle, stdout=file_handle) class SlurmEnvironment(Environment):

    @classmethod
    def is_usable(cls):
        return shutil.which("sbatch") is not None

    @classmethod
    def list_up_jobs(cls, user=None):
        # Taken from clusterlib (https://github.com/clusterlib/clusterlib)
        command = ["squeue", "--noheader", "-o", "%j"]
        if user is not None:
            command.extend(["-u", user])

        try:
            with open(os.devnull, 'w') as shutup:
                out = subprocess.check_output(command, stderr=shutup)
                return out.decode('utf-8').splitlines()
        except OSError:
            # OSError is raised if the program is not installed
            return None

    def __init__(self, serializer=Serializer(), time="1:00:00", memory=4000,
                 partition=None, n_proc=None, gpu=None,
                 shell_script="#!/bin/bash", fail_fast=True, other_flags=None,
                 other_options=None):
        super(SlurmEnvironment, self).__init__(fail_fast)
        self.serializer = serializer
        self.time = time
        self.memory = memory
        self.shell_script = shell_script
        self.partition = partition
        self.n_proc = n_proc
        self.gpu = gpu
        self.other_flags = [] if other_flags is None else other_flags
        self.other_options = {} if other_options is None else other_options

    def __repr__(self):
        return "{cls}(serializer={serializer}, time={time}, memory={memory}, " \
               "partition={partition}, n_proc={n_proc}, gpu={gpu}," \
               " shell_script={shell}, fail_fast={fail_fast}, " \
               "other_flags={other_flags}, other_options={other_options})" \
               "".format(cls=self.__class__.__name__,
                         serializer=repr(self.serializer),
                         time=repr(self.time),
                         memory=repr(self.memory),
                         partition=repr(self.partition),
                         n_proc=repr(self.n_proc),
                         gpu=repr(self.gpu),
                         shell=self.shell_script,
                         fail_fast=repr(self.fail_fast),
                         other_flags=repr(self.other_flags),
                         other_options=repr(self.other_options))

    def issue(self, lazy_computation):
        # Making Slurm command
        comp_name = lazy_computation.comp_name
        log_folder = lazy_computation.storage.get_log_folder()
        log_prefix = os.path.join(log_folder, comp_name)

        env_params = {"time": self.time, "memory": self.memory,
                      "partition": self.partition, "n_proc": self.n_proc,
                      "gpu": self.gpu, "other_flags": self.other_flags,
                      "other_options": self.other_options}

        env_params = self._customize(lazy_computation, env_params)

        slurm_cmd = ["sbatch", "--job-name={}".format(comp_name),
                     "--time={}".format(env_params["time"]),
                     "--mem={}".format(env_params["memory"]),
                     "--output={}.%j.txt".format(log_prefix),
                     ]

        if self.partition is not None:
            slurm_cmd.append("--partition={}".format(env_params["partition"]))

        if self.n_proc is not None:
            slurm_cmd.append("--cpus-per-task={}".format(env_params["n_proc"]))

        if self.gpu is not None:
            slurm_cmd.append("--gres=gpu:{}".format(env_params["gpu"]))

        slurm_cmd.extend(env_params["other_flags"])

        for flag, value in env_params["other_options"].items():
            slurm_cmd.append("{}={}".format(flag, value))

        # Making computation command
        cmd_as_tuple = self.serializer.serialize_and_script(lazy_computation)
        str_cmd = " ".join([escape(s) for s in cmd_as_tuple])
        whole_cmd = "{shell}\n{cmd}".format(shell=self.shell_script,
                                            cmd=str_cmd)

        logger = logging.getLogger("clustertools")
        logger.debug(slurm_cmd)
        logger.debug(whole_cmd)

        # Running everything
        # Slurm commands are launched through shell.
        # 1. We can either dump a bash program and then create a subprocess
        #    to launch it
        # 2. Or we can pipe directly both part. This is the option we take
        #    here.
        with subprocess.Popen(slurm_cmd, stdin=subprocess.PIPE,
                              stderr=subprocess.PIPE,
                              stdout=subprocess.PIPE) as process:
            stdout, stderr = process.communicate(whole_cmd.encode("utf-8"))
            stdout = stdout.decode("utf-8")
            stderr = stderr.decode("utf-8")
            logger.debug(stdout)

            if len(stderr) > 0:
                logger.error(stderr)
            if process.returncode != 0:
                logger.error("Return code: {}".format(process.returncode))
                raise subprocess.CalledProcessError(process.returncode,
                                                    cmd=slurm_cmd,
                                                    output=stdout,
                                                    stderr=stderr) class Serializer(object):

    def __repr__(self):
        return "{cls}()".format(cls=self.__class__.__name__)

    def serialize(self, lazy_computation):
        """Return a unserializable string"""
        return dill.dumps(lazy_computation, -1)

    def deserialize(self, serlialized):
        """Return the object represented by the serialized string"""
        return dill.loads(serlialized)

    def serialize_and_script(self, lazy_computation):
        """Return a script to run the lazy_computation"""
        serialized = self.serialize(lazy_computation)
        return [sys.executable, '-c',
                'from {mod} import {cls};'
                '{repr}({serialized})'
                ''.format(mod=__name__,
                          cls=self.__class__.__name__,
                          repr=repr(self),
                          serialized=repr(serialized))]

    def deserialize_and_run(self, serialized):
        lazy_computation = self.deserialize(serialized)
        lazy_computation()

    def __call__(self, serialized):
        self.deserialize_and_run(serialized) class FileSerializer(Serializer):

    def serialize(self, lazy_computation):
        fname = "{}-{}.pkl".format(lazy_computation.comp_name, str(epoch()))
        fpath = lazy_computation.storage.get_messy_path(fname)
        with open(fpath, "wb") as hdl:
            dill.dump(lazy_computation, hdl, -1)
        return fpath

    def deserialize(self, serlialized):
        with open(serlialized, "rb") as hdl:
            lazy_computation = dill.load(hdl)
        try:
            os.remove(serlialized)
        except IOError:
            pass
        return lazy_computation class DebugEnvironment(Environment):
    """
    `DebugEnvironment`
    ==================
    An `Environment` that do not run code. It only prints the messages.
    """

    @classmethod
    def is_usable(cls):
        return True

    def __init__(self, print_all_parameters=False, fail_fast=True):
        super().__init__(fail_fast)
        self.print_all_parameters = print_all_parameters

    def log(self, msg, *args, **kwargs):
        logger = logging.getLogger("clustertools")
        logger.debug(msg, *args, **kwargs)

    def run(self, experiment, start=0, capacity=None):
        with catch_logging():
            self.log("From start={}, with capacity={}, "
                     "running experiment '{}': "
                     "{}".format(start, capacity, experiment.exp_name,
                                 repr(experiment)))
            if self.print_all_parameters:
                print("Parameters are:")
                for x in experiment.parameter_set:
                    print(x)
                print()
            super().run(experiment, start, capacity)

    def issue(self, lazy_computation):
        self.log("Pseudo issuing computation '{}': {}"
                 "".format(lazy_computation.comp_name, repr(lazy_computation)))

    def __repr__(self):
        return "{cls}(print_all_parameters={print_all_parameters}, " \
               "fail_fast={fail_fast})"\
               "".format(cls=self.__class__.__name__,
                         print_all_parameters=repr(self.print_all_parameters),
                         fail_fast=repr(self.fail_fast))

    def create_session(self, experiment):
        return DebugSession(self).init(len(experiment), experiment.storage) class PickleStorage(Storage):
    """Databases are folders. Each record is an individual pickle file.
    Implements a backup mechanism for results to prevent long writes from
    being interrupted
    """

    @classmethod
    def _save(cls, stuff, fpath):
        with open(fpath, "wb") as hdl:
            pickle.dump(stuff, hdl, -1)

    @classmethod
    def _raw_load(cls, fpath):
        with open(fpath, "rb") as hdl:
            return pickle.load(hdl)

    @classmethod
    def _load(cls, fpath):
        try:
            rtn = cls._raw_load(fpath)
        except EOFError:
            logger = logging.getLogger("clustertools.storage")
            logger.error("End of file encountered in '{}'".format(fpath))
            return {}
        except Exception as exception:
            logger = logging.getLogger("clustertools.storage")
            logger.exception("Error while loading file '{}'. Reason: {}"
                             "".format(fpath, repr(exception)))
            return {}
        return rtn

    @property
    def _parameter_set_path(self):
        return os.path.join(self.folder, "parameter_set.pkl")

    def save_parameter_set(self, parameter_set):
        self._save(parameter_set, self._parameter_set_path)

    def load_parameter_set(self):
        if not os.path.exists(self._parameter_set_path):
            raise FileNotFoundError("Parameter set file '{}' does not exist. Make sure to run the experiment script "
                                    "at least once to save the parameters.".format(self._parameter_set_path))
        return self._load(self._parameter_set_path)

    def init(self):
        super(PickleStorage, self).init()
        for folder in self._get_notif_db(), self._get_result_db(), \
                self._get_tmp_folder():
            if not os.path.exists(folder):
                os.makedirs(folder)
        return self

    def _get_notif_db(self):
        return os.path.join(self.folder, "notifications")

    def _get_result_db(self):
        return os.path.join(self.folder, "results")

    def _get_tmp_folder(self):
        return os.path.join(self.folder, "temp")

    # |--------------------------- Notifications ----------------------------> #

    def update_state(self, state):
        fpath = os.path.join(self._get_notif_db(), "%s.pkl" % state.comp_name)
        self._save(state, fpath)
        return state

    def load_states(self):
        from .state import State
        res = []
        for fpath in glob.glob(os.path.join(self._get_notif_db(), "*.pkl")):
            loaded = self._load(fpath)
            try:
                if isinstance(loaded, State) or len(loaded) > 0:
                    res.append(loaded)
            except (AttributeError, TypeError, ValueError):
                # If `loaded` has no length (should be a type error)
                pass
        return res

    # |---------------------------- Results -----------------------------> #

    def _tmp_path(self, comp_name):
        return os.path.join(self._get_tmp_folder(), "%s.pkl" % comp_name)

    def _bc_path(self, comp_name):
        return os.path.join(self._get_tmp_folder(), "%s.bc.pkl" % comp_name)

    def _result_path(self, comp_name):
        return os.path.join(self._get_result_db(), "%s.pkl" % comp_name)

    def _save_r_dict(self, comp_name, r_dict):
        # Build paths
        tmp_path = self._tmp_path(comp_name)
        bc_path = self._bc_path(comp_name)
        fpath = self._result_path(comp_name)
        # Write to disk in temp
        self._save(r_dict, tmp_path)
        # Back up old results
        if os.path.exists(fpath):
            shutil.move(fpath, bc_path)
        # Place new file at the right place
        shutil.move(tmp_path, fpath)

    def restore_back_up(self, comp_name):
        bc_path = self._bc_path(comp_name)
        fpath = self._result_path(comp_name)
        shutil.move(bc_path, fpath)

    def _load_r_dict(self, comp_name):
        fpath = self._result_path(comp_name)
        if os.path.exists(fpath):
            return self._load(fpath)
        return {}

    def _load_r_dicts(self):
        """load and return all the proxy results"""
        r_dict = {}
        for fpath in glob.glob(os.path.join(self._get_result_db(), "*.pkl")):
            r_dict.update(self._load(fpath))
        return r_dict def purge(exp_name=__EXP_NAME__, storage_factory=IntrospectStorage):
    storage = storage_factory(experiment_name=exp_name)
    try:
        storage.delete()
    except OSError:
        pass
    return storage def prep(exp_name=__EXP_NAME__, storage_factory=IntrospectStorage):
    storage = purge(exp_name, storage_factory)
    storage.init() def pickle_prep(exp_name=__EXP_NAME__):
    prep(exp_name, PickleStorage) def pickle_purge(exp_name=__EXP_NAME__):
    purge(exp_name, PickleStorage) __EXP_NAME__ = "ClustertoolsTest" class IntrospectStorage(Storage):
    """
    This `Storage` does not store anything on disk but rather keeps track
    of everything that has happened.

    Attributes
    ----------
    state_history: mapping: computation_name -> list of `State`
        The list of state of that computation (most recent at the end)
    result_history: mapping: computation_name -> list of saved results
        The list of results of that computation (most recent at the end)
    """
    def __init__(self, experiment_name, architecture=Architecture()):
        super(IntrospectStorage, self).__init__(experiment_name, architecture)
        self.state_history = defaultdict(list)
        self.result_history = defaultdict(list)
        self.parameter_set_history = defaultdict(list)

    def update_state(self, state):
        self.state_history[state.comp_name].append(state)
        return state

    def load_states(self):
        return [states[-1] for states in self.state_history.values()]

    def _save_r_dict(self, comp_name, r_dict):
        self.result_history[comp_name].append(r_dict)

    def _load_r_dict(self, comp_name):
        return self.result_history[comp_name][-1]

    def _load_r_dicts(self):
        return {comp_name: history[-1] for comp_name, history
                in self.result_history.items()}

    def save_parameter_set(self, parameter_set):
        self.parameter_set_history[self.exp_name].append(parameter_set) class TestComputation(Computation):
    def __init__(self, exp_name=__EXP_NAME__, comp_name="TestComp",
                 context="n/a", storage_factory=IntrospectStorage):
        super(TestComputation, self).__init__(exp_name=exp_name,
                                              comp_name=comp_name,
                                              context=context,
                                              storage_factory=storage_factory)

    def run(self, result, x1, x2, **ignored):
        result["mult"] = x1 * x2 def with_setup_(setup=None, teardown=None):
    """Decorator like `with_setup` of nosetest but which can be applied to any
    function"""
    def decorated(function):
        def app(*args, **kwargs):
            if setup:
                setup()
            try:
                function(*args, **kwargs)
            finally:
                if teardown:
                    teardown()
        return app
    return decorated def skip_if_usuable(environment_cls):
    def skip_or_not():
        if not environment_cls.is_usable():
            raise SkipTest("Environment '{}' is not usable here"
                           "".format(repr(environment_cls)))

    return with_setup(skip_or_not)
import shutil
from functools import partial
from nose.tools import assert_equal, assert_in, assert_less, assert_raises, \
    with_setup, assert_true, assert_false
from nose.tools import assert_is_instance
from nose.tools import assert_is_none
from clustertools import ParameterSet, Experiment
from clustertools.environment import InSituEnvironment, \
    BashEnvironment, SlurmEnvironment, Serializer, FileSerializer, \
    DebugEnvironment
from clustertools.storage import PickleStorage
from .util_test import purge, prep, pickle_prep, pickle_purge, \
    __EXP_NAME__, IntrospectStorage, TestComputation, with_setup_, \
    skip_if_usuable
# -*- coding: utf-8 -*-




__author__ = "Begon Jean-Michel <jm.begon@gmail.com>"
__copyright__ = "3-clause BSD License"


# ---------------------------------------------------------------------- Session
# Assumes InSituEnvironment works correctly

@with_setup(prep, purge)
def test_session():
    parameter_set = ParameterSet()
    parameter_set.add_parameters(x1=range(3), x2=range(3))
    experiment = Experiment(__EXP_NAME__, parameter_set, TestComputation,
                            IntrospectStorage)

    env = InSituEnvironment(fail_fast=True)
    session = env.create_session(experiment)
    assert_false(session.is_open())
    assert_raises(ValueError, partial(session.run, TestComputation()))
    with session:
        for lazy_computation in experiment.yield_computations():
            session.run(lazy_computation)
        assert_equal(session.n_launch, 9)

    assert_false(session.is_open())
    assert_raises(ValueError, partial(session.run, TestComputation()))


# ----------------------------------------------------- Environment class method
def test_usability():
    assert_equal(InSituEnvironment.is_usable(), True)
    assert_equal(BashEnvironment.is_usable(), True)
    assert_equal(SlurmEnvironment.is_usable(),
                 shutil.which("scontrol") is not None)


def test_list_up_jobs():
    assert_is_none(InSituEnvironment.list_up_jobs())
    assert_is_none(BashEnvironment.list_up_jobs())

    if not SlurmEnvironment.is_usable():
        assert_is_none(SlurmEnvironment.list_up_jobs())

    # TODO else case


# ---------------------------------------------------------- Environment methods
def test_in_situ_autorefresh():
    env = InSituEnvironment(fail_fast=True)
    assert_true(env.auto_refresh)


def test_context_is_str():
    for env_cls in DebugEnvironment, InSituEnvironment, BashEnvironment, \
                   SlurmEnvironment:
        env = env_cls()
        assert_is_instance(env.context(), str)


# ------------------------------------------------- Environment generated script
@with_setup_(prep, purge)
def script_evaluation(environment):
    print(repr(environment))  # In case of error, prints the type of environment
    lazy_computation = TestComputation().lazyfy(x1=1, x2=2)
    # TODO check computation


def test_bash_script():
    script_evaluation(BashEnvironment())



@skip_if_usuable(SlurmEnvironment)