info.pancancer.arch3.jobGenerator.JobGenerator.java Source code

Java tutorial

Introduction

Here is the source code for info.pancancer.arch3.jobGenerator.JobGenerator.java

Source

/*
 *     Consonance - workflow software for multiple clouds
 *     Copyright (C) 2016 OICR
 *
 *     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/>.
 *
 */

package info.pancancer.arch3.jobGenerator;

import com.google.common.base.Joiner;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.MessageProperties;
import info.pancancer.arch3.Base;
import info.pancancer.arch3.beans.Job;
import info.pancancer.arch3.beans.Order;
import info.pancancer.arch3.beans.Provision;
import info.pancancer.arch3.persistence.PostgreSQL;
import info.pancancer.arch3.utils.Constants;
import info.pancancer.arch3.utils.IniFile;
import info.pancancer.arch3.utils.Utilities;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.TimeoutException;
import joptsimple.ArgumentAcceptingOptionSpec;
import joptsimple.OptionSpecBuilder;
import org.apache.commons.configuration.HierarchicalINIConfiguration;
import org.apache.tools.ant.DirectoryScanner;

/**
 * Created by boconnor on 15-04-18.
 *
 * takes a --config option pointed to a config .json file
 */
public class JobGenerator extends Base {

    // variables
    private HierarchicalINIConfiguration settings = null;
    private Channel jchannel = null;
    private String queueName = null;
    private int currIterations = 0;

    private final ArgumentAcceptingOptionSpec<String> iniDirSpec;
    private final ArgumentAcceptingOptionSpec<String> workflowNameSpec;
    private final ArgumentAcceptingOptionSpec<String> workflowVersionSpec;
    private final ArgumentAcceptingOptionSpec<String> workflowPathSpec;
    private final ArgumentAcceptingOptionSpec<Integer> totalJobSpec;
    private final OptionSpecBuilder forceSpec;

    public static void main(String[] argv) throws IOException {
        JobGenerator jg = new JobGenerator(argv);
        jg.log.info("MASTER FINISHED, EXITING!");
    }

    public JobGenerator(String[] argv) throws IOException {
        super();

        this.iniDirSpec = super.parser.accepts("ini-dir", "schedule a batch of ini files from this directory")
                .withOptionalArg().ofType(String.class);
        this.workflowNameSpec = super.parser.accepts("workflow-name", "track the name of workflows")
                .withOptionalArg().ofType(String.class).required();
        this.workflowVersionSpec = super.parser.accepts("workflow-version", "track the version of workflows")
                .withOptionalArg().ofType(String.class).required();
        this.workflowPathSpec = super.parser
                .accepts("workflow-path", "Schedule workflows at this path on the container host").withOptionalArg()
                .ofType(String.class).required();
        this.totalJobSpec = super.parser.accepts("total-jobs", "Schedule a specific number of test workflows")
                .requiredUnless(iniDirSpec, this.endlessSpec).withRequiredArg().ofType(Integer.class)
                .defaultsTo(Integer.MAX_VALUE);
        this.forceSpec = super.parser.accepts("force", "Force job generation even if hashing is activated");

        parseOptions(argv);

        String iniDir = options.has(iniDirSpec) ? options.valueOf(iniDirSpec) : null;
        String workflowName = options.valueOf(workflowNameSpec);
        String workflowVersion = options.valueOf(workflowVersionSpec);
        String workflowPath = options.valueOf(workflowPathSpec);

        // UTILS OBJECT
        settings = Utilities.parseConfig(configFile);

        // CONFIG
        queueName = settings.getString(Constants.RABBIT_QUEUE_NAME);
        log.info("queue name: " + queueName);
        try {
            // SETUP QUEUE
            this.jchannel = Utilities.setupQueue(settings, queueName + "_orders");
        } catch (TimeoutException ex) {
            throw new RuntimeException(ex);
        }

        if (options.has(iniDirSpec)) {
            // read an array of files
            log.info("scanning: " + iniDir);
            DirectoryScanner scanner = new DirectoryScanner();
            scanner.setIncludes(new String[] { "**/*.ini" });
            scanner.setBasedir(iniDir);
            scanner.setCaseSensitive(false);
            scanner.scan();
            String[] files = scanner.getIncludedFiles();

            // LOOP, ADDING JOBS EVERY FEW MINUTES
            for (String file : files) {
                generateAndQueueJob(iniDir + "/" + file, workflowName, workflowVersion, workflowPath);
            }
        } else if (options.has(endlessSpec) || options.has(totalJobSpec)) {
            // limit
            log.info("running in test mode");
            boolean endless = options.has(endlessSpec);
            int limit = options.valueOf(totalJobSpec);
            for (int i = 0; endless || i < limit; i++) {
                generateAndQueueJob(null, workflowName, workflowVersion, workflowPath);
            }
        }

        try {

            jchannel.getConnection().close(FIVE_SECOND_IN_MILLISECONDS);

        } catch (IOException ex) {
            log.error(ex.toString());
        }

    }

    private void generateAndQueueJob(String iniFile, String workflowName, String workflowVersion,
            String workflowPath) {
        // keep track of the iterations
        currIterations++;
        log.info("\ngenerating new jobs, iteration " + currIterations + "\n");
        // TODO: this is fake, in a real program this is being read from JSONL file or web service
        // check to see if new results are available and/or if the work queue is empty
        Order o = generateNewJob(iniFile, workflowName, workflowVersion, workflowPath);
        // enqueue new job
        if (o != null) {
            enqueueNewJobs(o.toJSON());
        }
        try {
            // pause
            Thread.sleep(ONE_SECOND_IN_MILLISECONDS);
        } catch (InterruptedException ex) {
            log.error(ex.getMessage(), ex);
        }
    }

    // PRIVATE

    // TODO: this will actually need to come from a file or web service
    private Order generateNewJob(String file, String workflowName, String workflowVersion, String workflowPath) {

        Map<String, String> iniFileEntries;
        if (file != null) {
            // TODO: this will come from a web service or file
            iniFileEntries = parseIniFile(file);
            for (Entry<String, String> e : iniFileEntries.entrySet()) {
                log.info("key: " + e.getKey() + " value: " + e.getValue());
            }
        } else {
            iniFileEntries = new LinkedHashMap<>();
            iniFileEntries.put("param1", "bar");
            iniFileEntries.put("param2", "foo");
        }

        String hashStr;
        PostgreSQL db = new PostgreSQL(settings);

        Joiner.MapJoiner mapJoiner = Joiner.on(',').withKeyValueSeparator("=");
        String[] arr = this.settings.getStringArray(Constants.JOB_GENERATOR_FILTER_KEYS_IN_HASH);
        Map<String, String> filteredIniFileEntries = iniFileEntries;
        if (arr.length > 0) {
            log.info("Using ini hash filter set: " + Arrays.toString(arr));
            Set<String> keys = new HashSet<>();
            Map<String, String> filteredMap = new LinkedHashMap<>();
            keys.addAll(Arrays.asList(arr));
            for (Entry<String, String> entry : iniFileEntries.entrySet()) {
                if (keys.contains(entry.getKey())) {
                    filteredMap.put(entry.getKey(), entry.getValue());
                }
            }
            filteredIniFileEntries = filteredMap;
        }
        // don't use a real hashcode here, they can duplicate
        // just use the value of the filtered map
        hashStr = String.valueOf(mapJoiner.join(filteredIniFileEntries));

        if (this.settings.getBoolean(Constants.JOB_GENERATOR_CHECK_JOB_HASH, Boolean.FALSE)) {
            boolean runPreviously = db.previouslyRun(hashStr);
            if (runPreviously) {
                if (this.options.has(this.forceSpec)) {
                    log.info("Forcing scheduling, but would have skipped file (null if testing) due to hash: "
                            + file);
                } else {
                    log.info("Skipping file (null if testing) due to hash: " + file);
                    return null;

                }
            }
        }

        int cores = DEFAULT_NUM_CORES;
        int memGb = DEFAULT_MEMORY;
        int storageGb = DEFAULT_DISKSPACE;
        ArrayList<String> a = new ArrayList<>();
        a.add("ansible_playbook_path");

        Order newOrder = new Order();
        newOrder.setJob(new Job(workflowName, workflowVersion, workflowPath, hashStr, iniFileEntries));
        newOrder.setProvision(new Provision(cores, memGb, storageGb, a));
        newOrder.getProvision().setJobUUID(newOrder.getJob().getUuid());

        return newOrder;

    }

    private Map<String, String> parseIniFile(String iniFile) {

        Map<String, String> iniHash = new LinkedHashMap<>();

        try {
            IniFile ini = new IniFile(iniFile);
            iniHash = ini.getEntries().get("no-section");

        } catch (IOException e) {
            log.error(e.toString());
        }

        return iniHash;
    }

    private void enqueueNewJobs(String job) {

        try {
            log.info("\nSENDING JOB:\n '" + job + "'\n" + this.jchannel + " \n");

            this.jchannel.basicPublish("", queueName + "_orders", MessageProperties.PERSISTENT_TEXT_PLAIN,
                    job.getBytes(StandardCharsets.UTF_8));
            jchannel.waitForConfirms();
        } catch (IOException | InterruptedException ex) {
            log.error(ex.toString());
        }

    }

}