dispatch.cc

/* Spell-QTL  Software suite for the QTL analysis of modern datasets.
 * Copyright (C) 2016,2017  Damien Leroux <damien.leroux@inra.fr>, Sylvain Jasson <sylvain.jasson@inra.fr>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <atomic>
#include <thread>
#include <cstdio>
#include "eigen.h"
#include "input.h"
#include "dispatch.h"
#include "inet.h"

extern "C" {
#include <sys/stat.h> 
#include <fcntl.h>
#include<signal.h>
#include<sys/types.h>    
}

size_t count_jobs(const std::string& name, bn_settings_t* settings) { return job_registry.find(name)->second.first(settings); }
job_type find_job(const std::string& name) { return job_registry.find(name)->second.second; }

struct worker_base {
    size_t m_slice_begin, m_slice_end;
    bn_settings_t* m_settings;
    job_type run_one;

    worker_base(bn_settings_t* settings, size_t slice_begin, size_t slice_end, const std::string& job)
        : m_slice_begin(slice_begin), m_slice_end(slice_end), m_settings(settings), run_one(find_job(job))
    {}

    void run_slice()
    {
        /*MSG_INFO("RUN SLICE " << m_slice_begin << ':' << m_slice_end);*/
        for (size_t i = m_slice_begin; i <= m_slice_end; ++i) {
            auto ret = run_one(m_settings, i);
            notify_one(i, ret);
        }
    }

    /* middleware */
    virtual void notify_one(size_t, bool) = 0;
    virtual void join() = 0;
};


struct worker_local_ssh : public worker_base {
    using worker_base::worker_base;
    void notify_one(size_t i, bool ok)
    {
        msg_handler_t::cout() << (ok ? "\fS " : "\fF ") << i << std::endl;
    }
    void join() {}
};


struct worker_local_sge : public worker_base {
    using worker_base::worker_base;
    void notify_one(size_t i, bool ok)
    {
        MSG_INFO("Opening FIFO for notification. " << m_settings->fifo_path.c_str());
        MSG_QUEUE_FLUSH();
        MSG_INFO("Sending notification.");
        MSG_QUEUE_FLUSH();

        inet_client client(m_settings->fifo_path);
        if (errno) {
            MSG_ERROR("An error occurred while trying to connect to the master process: " << strerror(errno), "");
            MSG_QUEUE_FLUSH();
        } else {
            std::string msg = SPELL_STRING((ok ? "\fS " : "\fF ") << i << std::endl);
            client.connect_and_send(msg);
            if (errno) {
                MSG_ERROR("An error occurred while trying to notify the master process: " << strerror(errno), "");
                MSG_QUEUE_FLUSH();
            }
        }
/*
        FILE* fifo = fopen(m_settings->fifo_path.c_str(), "a");
        if (!fifo || errno) {
            MSG_ERROR("An error occurred while trying to notify the master process: " << strerror(errno), "");
        }
        std::string msg = SPELL_STRING((ok ? "\fS " : "\fF ") << i << std::endl);
        fwrite(msg.c_str(), msg.size(), 1, fifo);
        if (ferror(fifo)) {
            MSG_ERROR("An error occurred while trying to notify the master process: " << strerror(errno), "");
        }
        fflush(fifo);
        fclose(fifo);
*/
        MSG_INFO("Done notifying.");
    }
    void join() {}
};



struct master_base {
    bn_settings_t* m_settings;
    std::vector<std::shared_ptr<worker_base>> m_workers;
    size_t m_n_jobs;
    size_t m_job_total;
    std::atomic<size_t> m_job_done_counter;
    std::atomic<size_t> m_job_ok_counter;
    std::atomic<size_t> m_job_failed_counter;
    std::string job_name;
    master_base(bn_settings_t* settings, size_t n_jobs, const std::string& job)
        : m_settings(settings)
        , m_workers()
        , m_n_jobs(n_jobs)
        , m_job_total(0)
        , m_job_done_counter(0)
        , m_job_ok_counter(0)
        , m_job_failed_counter(0)
        , job_name(job)
    {}

    void run()
    {
        double step;
        m_job_total = count_jobs(job_name, m_settings);
        size_t n_slots;
        MSG_DEBUG("Have " << m_n_jobs << " slots for " << m_job_total << " jobs to do.");
        if (m_n_jobs >= m_job_total) {
            n_slots = m_job_total;
            step = 0;
        } else {
            n_slots = m_n_jobs;
            step = m_job_total / (double) n_slots;
        }
        /*MSG_DEBUG("step=" << step);*/

        m_workers.clear();
        m_workers.reserve(n_slots);

        double start = 0;
        for (size_t i = 0; i < n_slots; ++i) {
            size_t end = (size_t) (start + step);
            if (end >= m_job_total) {
                end = m_job_total - 1;
            }
            /*MSG_DEBUG("Spawning worker for " << ((size_t) start) << ':' << end);*/
            m_workers.emplace_back(spawn_worker(i, (size_t) start, end, job_name));
            start += step + 1;
        }
    }

    bool all_good() const { return m_job_failed_counter == 0; }

    virtual std::shared_ptr<worker_base> spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job) = 0;
    virtual void wait_for_jobs() = 0;

    void update_job(bool successful)
    {
        if (successful) {
            ++m_job_ok_counter;
        } else {
            ++m_job_failed_counter;
        }
        ++m_job_done_counter;
        msg_handler_t::cout() << "\033[100D\033[KProgress: " << YELLOW << m_job_done_counter << NORMAL << '/' << YELLOW << m_job_total << NORMAL << " (" << m_job_ok_counter << " good, " << m_job_failed_counter << " failed)" << std::flush;
    }
};


/**************************************
 * MULTI-THREADING
 */

struct master_none;
struct worker_none;

struct master_none : public master_base {
    master_none(bn_settings_t* settings, const std::string& job)
        : master_base(settings, 1, job)
    {}

    void wait_for_jobs();

    void notify(size_t /*i*/, bool ok)
    {
        update_job(ok);
    }

    std::shared_ptr<worker_base> spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job);
};


struct worker_none : public worker_base {
    master_none* m_master;
    worker_none(master_none* master, size_t slice_begin, size_t slice_end, const std::string& job)
        : worker_base(master->m_settings, slice_begin, slice_end, job)
        , m_master(master)
    {}

    void notify_one(size_t i, bool ok)
    {
        m_master->notify(i, ok);
    }

    void join() { run_slice(); }
};


void master_none::wait_for_jobs()
{
    for (auto& t: m_workers) { t->join(); }
}

std::shared_ptr<worker_base> master_none::spawn_worker(size_t, size_t slice_start ,size_t slice_end, const std::string& job)
{
    return std::dynamic_pointer_cast<worker_base>(std::make_shared<worker_none>(this, slice_start, slice_end, job));
}


/**************************************
 * MULTI-THREADING
 */

struct master_thread;
struct worker_thread;

struct master_thread : public master_base {
    std::mutex m_notify_lock;
    master_thread(bn_settings_t* settings, const std::string& job)
        : master_base(settings, settings->n_threads, job)
        , m_notify_lock()
    {}

    std::shared_ptr<worker_base> spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job);

    void wait_for_jobs();

    void notify(size_t /*i*/, bool ok)
    {
        std::lock_guard<std::mutex> lock(m_notify_lock);
        update_job(ok);
    }
};

struct worker_thread : public worker_base {
    master_thread* m_master;
    std::thread m_thread;
    worker_thread(master_thread* master, size_t slice_begin, size_t slice_end, const std::string& job)
        : worker_base(master->m_settings, slice_begin, slice_end, job)
        , m_master(master)
        , m_thread([&] () { run_slice(); })
    {}

    void notify_one(size_t i, bool ok)
    {
        m_master->notify(i, ok);
    }

    void join()
    {
        m_thread.join();
    }
};


void master_thread::wait_for_jobs()
{
    for (auto& t: m_workers) { t->join(); }
}


std::shared_ptr<worker_base> master_thread::spawn_worker(size_t, size_t slice_start ,size_t slice_end, const std::string& job)
{
    return std::dynamic_pointer_cast<worker_base>(std::make_shared<worker_thread>(this, slice_start, slice_end, job));
}


/**************************************
 * SSH
 */

struct master_ssh;
struct worker_ssh;

struct master_ssh : public master_base {
    master_ssh(bn_settings_t* settings, const std::string& job)
        : master_base(settings, settings->ssh_hosts.size(), job)
    {}

    std::shared_ptr<worker_base> spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job);

    void wait_for_jobs();

    void notify(size_t /*i*/, bool ok)
    {
        update_job(ok);
    }
};

struct worker_ssh : public worker_base {
    size_t m_host;
    master_ssh* m_master;
    std::string m_job_name;
    std::thread m_thread;
    worker_ssh(master_ssh* master, size_t i, size_t slice_begin, size_t slice_end, const std::string& job)
        : worker_base(master->m_settings, slice_begin, slice_end, job)
        , m_host(i)
        , m_master(master)
        , m_job_name(job)
        , m_thread([this] () { run_and_monitor(); })
    {}

    void notify_one(size_t i, bool ok)
    {
        m_master->notify(i, ok);
    }

    void join()
    {
        m_thread.join();
    }

    void run_and_monitor()
    {
        std::stringstream cmd;
        cmd << "ssh " << m_settings->ssh_hosts[m_host];
        for (const std::string& a: m_settings->command_line) {
            cmd << ' ' << a;
        }
        cmd << " -J " << m_job_name << ' ' << m_slice_begin << ' ' << m_slice_end;
        MSG_DEBUG("RUNNING " << cmd.str());
        FILE* remote = popen(cmd.str().c_str(), "r");
        char buf[64];
        while (!feof(remote)) {
            char* s = fgets(buf, (sizeof buf) - 1, remote);
            if (s && s[0] == '\f' && (s[1] == 'S' || s[1] == 'F')) {
                notify_one(atoi(s + 3), s[1] == 'S');
            }
        }
    }
};


void master_ssh::wait_for_jobs()
{
    for (auto& t: m_workers) { t->join(); }
}


std::shared_ptr<worker_base> master_ssh::spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job)
{
    return std::dynamic_pointer_cast<worker_base>(std::make_shared<worker_ssh>(this, i, slice_start, slice_end, job));
}



/**************************************
 * SGE
 */

struct master_sge;
struct worker_sge;

struct master_sge : public master_base, inet_server {
    std::mutex wait;
    size_t jobs_done;
    size_t jobs_total;

    master_sge(bn_settings_t* settings, const std::string& job)
        : master_base(settings, settings->n_threads, job), inet_server(56000, 65535)
        , jobs_done(0)
        , jobs_total(count_jobs(job, settings))
    {
        wait_for_server_to_run();
        m_settings->fifo_path = address();
    }

    std::shared_ptr<worker_base> spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job) override;

    void wait_for_jobs() override { wait.lock(); wait.unlock(); }

    void notify(size_t /*i*/, bool ok)
    {
        ++jobs_done;
        /*MSG_DEBUG("Job #" << i << (ok ? " successful" : " failed") << " (" << jobs_done << '/' << jobs_total << ')');*/
        update_job(ok);
        if (jobs_done >= jobs_total) {
            wait.unlock();
        }
    }

    bool on_client_message(const std::string& buf) override
    {
        if (buf.size() >= 3) {
            if (buf[0] == '\f') {
                if (buf[1] == 'F' || buf[1] == 'S') {
                    notify(atoi(buf.c_str() + 3), buf[1] == 'S');
                }
            }
        }
        return jobs_done < jobs_total;
    }

    bool on_client_error(int) override { return jobs_done < jobs_total; }

    void on_start() override { wait.lock(); }
    void on_stop() override { wait.unlock(); }

};

struct worker_sge : public worker_base {
    size_t m_host;
    master_sge* m_master;
    std::string m_job_name;
    worker_sge(master_sge* master, size_t i, size_t slice_begin, size_t slice_end, const std::string& job)
        : worker_base(master->m_settings, slice_begin, slice_end, job)
        , m_host(i)
        , m_master(master)
        , m_job_name(job)
    { schedule(); }

    void notify_one(size_t i, bool ok)
    {
        m_master->notify(i, ok);
    }

    void join() {}

    void schedule()
    {
        char wdbuf[PATH_MAX];
        const char* wd = getcwd(wdbuf, PATH_MAX);
        std::stringstream cmd;
        cmd << "cd " << wd << " &&";
        for (const std::string& a: m_settings->command_line) {
            cmd << ' ' << a;
        }
        cmd << " -J " << m_job_name << ' ' << m_slice_begin << ' ' << m_slice_end << " -F " << m_settings->fifo_path;
        /*MSG_DEBUG("SCHEDULING JOB " << cmd.str());*/
        std::string qsub = SPELL_STRING("/usr/bin/env qsub -N " << m_settings->prg_name << '.' << m_job_name << " " << m_settings->qsub_opts << " > /dev/null");
        FILE* remote = popen(qsub.c_str(), "w");
        std::string cmdstr = cmd.str();
        /*const char* cmdstr = cmd.str().c_str();*/
        fwrite(cmdstr.c_str(), cmdstr.size(), 1, remote);
        fclose(remote);
    }
};


std::shared_ptr<worker_base> master_sge::spawn_worker(size_t i, size_t slice_start ,size_t slice_end, const std::string& job)
{
    return std::dynamic_pointer_cast<worker_base>(std::make_shared<worker_sge>(this, i, slice_start, slice_end, job));
}


bool do_the_job(bn_settings_t* settings, std::string job)
{
    bool all_good = false;
    if (settings->is_master()) {
        msg_handler_t::cout() << std::endl << GREEN << "Running jobs " << job << NORMAL << std::endl;
        if (settings->scheme == JDS_MT) {
            master_thread mt(settings, job);
            mt.run();
            mt.wait_for_jobs();
            all_good = mt.all_good();
        } else if (settings->scheme == JDS_SSH) {
            master_ssh mssh(settings, job);
            mssh.run();
            mssh.wait_for_jobs();
            all_good = mssh.all_good();
        } else if (settings->scheme == JDS_SGE) {
            master_sge msge(settings, job);
            msge.run();
            msge.wait_for_jobs();
            all_good = msge.all_good();
        } else {
            master_none mnone(settings, job);
            mnone.run();
            mnone.wait_for_jobs();
            all_good = mnone.all_good();
        }
        if (!all_good) {
            msg_handler_t::cout() << std::endl;
            MSG_ERROR("At least one job failed. Aborting.", "Check error messages in job logs.");
        }
        return all_good;
    } else {
        if (settings->scheme == JDS_SSH) {
            worker_local_ssh wl(settings, settings->job_start, settings->job_end, job);
            wl.run_slice();
        } else if (settings->scheme == JDS_SGE) {
            worker_local_sge wl(settings, settings->job_start, settings->job_end, job);
            wl.run_slice();
        }
        return true;
    }
}

#include "io/output_impl.h"