Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.