g5k_prepare_bench_flops

#!/usr/bin/env python

import sys, os
from execo import Remote, Put, Get, SequentialActions, Report
from execo_g5k import get_g5k_clusters, get_cluster_site, OarSubmission, g5k_configuration
from execo_engine import logger, ParamSweeper
from g5k_cluster_engine import g5k_cluster_engine, worker_log
from os.path import join as pjoin
from os.path import exists as pexists

from common import *

class g5k_prepare_bench_flops(g5k_cluster_engine):

    def __init__(self):
        super(g5k_prepare_bench_flops, self).__init__()
        self.options_parser.set_usage("usage: %prog <comma separated list of clusters>")
        self.options_parser.set_description("precompile package for benching flops")
        self.options_parser.add_option("-o", dest = "oar_options", help = "oar reservation options", default = None)
        self.options_parser.add_option("-w", dest = "walltime", help = "walltime of compilation jobs", type = "string", default = "8:0:0")
        self.options_parser.add_option("-T", dest = "testing", action = "store_true", default = False,
                                       help = "use api branch testing")
        self.options_parser.add_argument("clusters", "comma separated list of clusters. ALL for all clusters")
        self.prepare_path = pjoin(self.engine_dir, "preparation")

    def init(self):
        if len(self.args) != 1:
            print "ERROR: missing argument"
            self.options_parser.print_help(file=sys.stderr)
            exit(1)
        if self.options.testing:
            g5k_configuration['api_additional_args'].append("branch=testing")
        try:
            os.makedirs(self.prepare_path)
        except:
            pass
        clusters = set(self.args[0].split(","))
        if "ALL" in clusters:
            clusters.remove("ALL")
            clusters.update(get_g5k_clusters())
        clusters_todo = [
            cluster
            for cluster in clusters
            if (not pexists(pjoin(self.prepare_path, prepared_archive("atlas", cluster)))
                or not pexists(pjoin(self.prepare_path, prepared_archive("openmpi", cluster)))
                or not pexists(pjoin(self.prepare_path, prepared_archive("hpl", cluster)))) ]
        self.sweeper = ParamSweeper(pjoin(self.result_dir, "prepare_params"),
                                    clusters_todo,
                                    save_sweeps = True)
        logger.info("cluster = %s" % (clusters_todo,))

    def get_clusters(self):
        return self.sweeper.get_remaining()

    def get_job(self, cluster):
        comb = self.sweeper.get_next(filtr = lambda r: filter(lambda comb: comb == cluster, r))
        if not comb:
            return None
        submission = OarSubmission(resources = "{cluster='%s'}/nodes=1" % (cluster,),
                                   walltime = self.options.walltime,
                                   name = "flopscompilworker",
                                   job_type = ["allow_classic_ssh"],
                                   additional_options = self.options.oar_options)
        return submission, comb

    def worker(self, cluster, site, comb, nodes, worker_index, oarsubmission, jobid):

        if jobid == None or len(nodes) == 0:
            self.sweeper.cancel(comb)
            return

        def prepare_package(package):
            worker_log.info("preparing package %s" % package)
            worker_log.info("copy %s files to nodes" % package)
            preparation = SequentialActions(
                [ Remote("mkdir -p " + node_working_dir,
                         nodes),
                  Put(nodes,
                      local_files = ([ pjoin(self.engine_dir, packages[package]["archive"]),
                                       pjoin(self.engine_dir, "node_prepare_" +  package) ] +
                                     [ pjoin(self.prepare_path, prepared_archive(dep, cluster)) for dep in packages[package]["deps"] ]),
                      remote_location = node_working_dir)])
            preparation.run()
            if not preparation.ok:
                worker_log.info("aborting, copy of %s files failed:\n%s" % (package, Report([preparation]).to_string()))
                self.sweeper.cancel(comb)
                return
            worker_log.info("compile %s" % package)
            compil = Remote(
                "%s/node_prepare_%s %s/%s %s %s %s > %s/%s.stdout" % (
                    node_working_dir,
                    package,
                    node_working_dir,
                    packages[package]["archive"],
                    packages[package]["extract_dir"],
                    prepared_archive(package, cluster),
                    " ".join([ prepared_archive(dep, cluster) for dep in packages[package]["deps"] ]),
                    node_working_dir,
                    prepared_archive(package, cluster)),
                nodes)
            compil.run()
            if not compil.ok:
                worker_log.info("%s compilation failed:\n%s" % (package, Report([compil]).to_string()))
                self.sweeper.cancel(comb)
            worker_log.info("retrieve result of %s compilation" % package)
            remote_files = [ pjoin(node_working_dir, prepared_archive(package, cluster) + ".stdout") ]
            if compil.ok:
                remote_files.append(pjoin(node_working_dir, prepared_archive(package, cluster)))
            retrieval = Get(
                nodes,
                remote_files = remote_files,
                local_location = self.prepare_path)
            retrieval.run()
            if not retrieval.ok:
                try:
                    os.unlink(pjoin(self.prepare_path, prepared_archive(package, cluster)))
                except:
                    pass
                worker_log.info("aborting, retrieval of %s compilation results failed:\n%s" % (package, Report([retrieval]).to_string()))
                return
            if compil.ok and retrieval.ok:
                worker_log.info("finished compilation of package %s" % package)
                self.sweeper.done(comb)

        for package in ["atlas", "openmpi", "hpl"]:
            if not pexists(pjoin(self.prepare_path, prepared_archive(package, cluster))):
                prepare_package(package)

if __name__ == "__main__":
    e = g5k_prepare_bench_flops()
    e.start()