SlurmJobSubmitter.java

package org.rabix.backend.local.slurm;


import org.apache.commons.io.FileUtils;
import org.rabix.bindings.BindingException;
import org.rabix.bindings.Bindings;
import org.rabix.bindings.BindingsFactory;
import org.rabix.bindings.model.Job;
import org.rabix.bindings.model.requirement.ResourceRequirement;
import org.rabix.common.helper.JSONHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;

public class SlurmJobSubmitter {

    private final static Logger logger = LoggerFactory.getLogger(SlurmJobSubmitter.class);
    final private static String jobPath = "rabix-backend-local/src/main/java/org/rabix/backend/local/slurm/transcribe.json";   // hardcoded command-line argument
    final private static String slurmBunnyJobPath = "job.json";

    public static void main(String[] args) {
        File jobFile = new File(jobPath);
        try {
            Job job = JSONHelper.readObject(FileUtils.readFileToString(jobFile), Job.class);
            Bindings bindings = BindingsFactory.create(job);

            String slurmJob = "#!/bin/sh\n";

            ResourceRequirement resourceRequirements = bindings.getResourceRequirement(job);
            String slurmDirective = getSlurmResourceRequirements(resourceRequirements);
            slurmJob += slurmDirective;

            String slurmCommand = "srun /usr/share/rabix-slurm-command-line/rabix -j " + slurmBunnyJobPath;
            slurmJob += slurmCommand;
            
            // creating bunny job file in the execution node(s), this must be done via sharedFileStorage      
            File slurmBunnyJobFile = new File(slurmBunnyJobPath);
            FileUtils.writeStringToFile(slurmBunnyJobFile, JSONHelper.writeObject(job));

            String slurmSchedulerJobPath = "slurm-job.sh";
            File slurmSchedulerJob = new File(slurmSchedulerJobPath);
            FileUtils.writeStringToFile(slurmSchedulerJob, slurmJob);

            // run "sbatch slurm-job.sh"
            // and check status via `squeue` command

        } catch (IOException e) {
            logger.error("Could not open job file");
            e.printStackTrace(System.err);
            System.exit(10);
        } catch (BindingException e) {
            logger.error("Failed to use Bindings", e);
            e.printStackTrace(System.err);
            System.exit(11);
        }
    }

    private static String getSlurmResourceRequirements(ResourceRequirement requirements){
        final String batchDirective = "#SBATCH";
        String directive = "";
        Long cpuMin = requirements.getCpuMin();
        Long memMin = requirements.getMemMinMB();
        if (cpuMin != null){
            directive += batchDirective + " --ntasks-per-node=" + Long.toString(cpuMin) + "\n";
        }
        if (memMin != null){
            directive += batchDirective + " --mem=" + Long.toString(memMin) + "\n";
        }
        return directive;
    }
}