| package edu.iu.helix.airavata; |
| |
| import edu.iu.helix.airavata.tasks.HelixTaskA; |
| import edu.iu.helix.airavata.tasks.HelixTaskB; |
| import edu.iu.helix.airavata.tasks.HelixTaskC; |
| import edu.iu.helix.airavata.tasks.HelixTaskD; |
| import edu.iu.helix.airavata.tasks.ssh.SSHKeyAuthentication; |
| import edu.iu.helix.airavata.tasks.ssh.SSHServerInfo; |
| import edu.iu.helix.airavata.tasks.ssh.SSHTask; |
| import edu.iu.helix.airavata.tasks.ssh.SSHTaskContext; |
| import org.apache.commons.io.IOUtils; |
| import org.apache.curator.framework.CuratorFramework; |
| import org.apache.helix.task.*; |
| import org.jboss.netty.util.internal.ThreadLocalRandom; |
| |
| import java.io.*; |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.UUID; |
| |
| /** |
| * Created by goshenoy on 6/21/17. |
| */ |
| public class HelixUtil { |
| |
| public static final String TASK_STATE_DEF = "Task"; |
| public static final String ZK_ADDRESS = "localhost:2199"; |
| public static final String CLUSTER_NAME = "HelixDemoCluster"; |
| |
| public static final String SSH_WORKFLOW = "SSH_Workflow"; |
| public static final String CREATE_DIR_TASK = "Task_CreateDir"; |
| public static final String COPY_PBS_TASK = "Task_CopyPBS"; |
| public static final String COPY_PY_TASK = "Task_CopyPY"; |
| public static final String QSUB_TASK = "Task_QSUB"; |
| |
| |
| public static final String USERNAME = "username"; |
| public static final String PRIVATE_KEY = "private_key"; |
| public static final String PUBLIC_KEY = "public_key"; |
| public static final String HOSTNAME = "hostname"; |
| public static final String PORT = "port"; |
| public static final String COMMAND = "command"; |
| public static final String SRC_PATH = "src_path"; |
| public static final String DEST_PATH = "dest_path"; |
| |
| |
| public enum DAGType { |
| TYPE_A, |
| TYPE_B, |
| TYPE_C, |
| SSH |
| } |
| |
| private static Workflow.Builder getWorkflowBuilder(DAGType dagType) throws Exception { |
| Workflow.Builder workflow = null; |
| |
| if (dagType.equals(DAGType.SSH)) { |
| if (!setWorkflowData()) { |
| throw new Exception("Failed to create zk data for SSH workflow"); |
| } |
| // create dir task - 1 task for job |
| List<TaskConfig> createDirTask = new ArrayList<TaskConfig>(); |
| createDirTask.add(new TaskConfig.Builder().setTaskId(CREATE_DIR_TASK).setCommand(SSHTask.TASK_COMMAND).build()); |
| |
| // copy files task - 2 tasks for job |
| List<TaskConfig> copyFilesTask = new ArrayList<TaskConfig>(); |
| copyFilesTask.add(new TaskConfig.Builder().setTaskId(COPY_PBS_TASK).setCommand(SSHTask.TASK_COMMAND).build()); |
| copyFilesTask.add(new TaskConfig.Builder().setTaskId(COPY_PY_TASK).setCommand(SSHTask.TASK_COMMAND).build()); |
| |
| // qsub task - 1 task for job |
| List<TaskConfig> qsubTask = new ArrayList<TaskConfig>(); |
| qsubTask.add(new TaskConfig.Builder().setTaskId(QSUB_TASK).setCommand(SSHTask.TASK_COMMAND).build()); |
| |
| // create-dir job config |
| JobConfig.Builder createDirJob = new JobConfig.Builder().addTaskConfigs(createDirTask).setMaxAttemptsPerTask(3); |
| |
| // copy-files job config |
| JobConfig.Builder copyFilesJob = new JobConfig.Builder().addTaskConfigs(copyFilesTask).setMaxAttemptsPerTask(3); |
| |
| // qsub job config |
| JobConfig.Builder qsubJob = new JobConfig.Builder().addTaskConfigs(qsubTask).setMaxAttemptsPerTask(1); |
| |
| // create workflow |
| workflow = new Workflow.Builder(SSH_WORKFLOW).setExpiry(0); |
| workflow.addJob("createDirJob", createDirJob); |
| workflow.addJob("copyFilesJob", copyFilesJob); |
| workflow.addJob("qsubJob", qsubJob); |
| } else { |
| // create task configs |
| List<TaskConfig> taskConfig1 = new ArrayList<TaskConfig>(); |
| List<TaskConfig> taskConfig2 = new ArrayList<TaskConfig>(); |
| List<TaskConfig> taskConfig3 = new ArrayList<TaskConfig>(); |
| List<TaskConfig> taskConfig4 = new ArrayList<TaskConfig>(); |
| taskConfig1.add(new TaskConfig.Builder().setTaskId("helix_task_a").setCommand(HelixTaskA.TASK_COMMAND).build()); |
| taskConfig2.add(new TaskConfig.Builder().setTaskId("helix_task_b").setCommand(HelixTaskB.TASK_COMMAND).build()); |
| taskConfig3.add(new TaskConfig.Builder().setTaskId("helix_task_c").setCommand(HelixTaskC.TASK_COMMAND).build()); |
| taskConfig4.add(new TaskConfig.Builder().setTaskId("helix_task_d").setCommand(HelixTaskD.TASK_COMMAND).build()); |
| |
| // create job configs |
| JobConfig.Builder jobConfig1 = new JobConfig.Builder().addTaskConfigs(taskConfig1).setMaxAttemptsPerTask(3); |
| JobConfig.Builder jobConfig2 = new JobConfig.Builder().addTaskConfigs(taskConfig2).setMaxAttemptsPerTask(3); |
| JobConfig.Builder jobConfig3 = new JobConfig.Builder().addTaskConfigs(taskConfig3).setMaxAttemptsPerTask(3); |
| JobConfig.Builder jobConfig4 = new JobConfig.Builder().addTaskConfigs(taskConfig4).setMaxAttemptsPerTask(3); |
| |
| // create workflow |
| workflow = new Workflow.Builder("helix_workflow").setExpiry(0); |
| workflow.addJob("helix_job_a", jobConfig1); |
| workflow.addJob("helix_job_b", jobConfig2); |
| workflow.addJob("helix_job_c", jobConfig3); |
| workflow.addJob("helix_job_d", jobConfig4); |
| } |
| |
| switch (dagType) { |
| case TYPE_A: |
| workflow.addParentChildDependency("helix_job_a", "helix_job_b"); |
| workflow.addParentChildDependency("helix_job_b", "helix_job_c"); |
| workflow.addParentChildDependency("helix_job_c", "helix_job_d"); |
| break; |
| |
| case TYPE_B: |
| workflow.addParentChildDependency("helix_job_a", "helix_job_c"); |
| workflow.addParentChildDependency("helix_job_c", "helix_job_d"); |
| workflow.addParentChildDependency("helix_job_d", "helix_job_b"); |
| break; |
| |
| case TYPE_C: |
| workflow.addParentChildDependency("helix_job_a", "helix_job_d"); |
| workflow.addParentChildDependency("helix_job_d", "helix_job_b"); |
| workflow.addParentChildDependency("helix_job_b", "helix_job_c"); |
| break; |
| |
| case SSH: |
| workflow.addParentChildDependency("createDirJob", "copyFilesJob"); |
| workflow.addParentChildDependency("copyFilesJob", "qsubJob"); |
| break; |
| } |
| |
| return workflow; |
| } |
| |
| public static Workflow getWorkflow(DAGType dagType) throws Exception { |
| Workflow.Builder workflowBuilder = getWorkflowBuilder(dagType); |
| return workflowBuilder.build(); |
| } |
| |
| private static String generateWorkflowName() { |
| return "workflow_" + ThreadLocalRandom.current().nextInt(9999); |
| } |
| |
| private static boolean setWorkflowData() { |
| try { |
| CuratorFramework curatorClient = ZkUtils.getCuratorClient(); |
| |
| SSHKeyAuthentication br2SshAuthentication = new SSHKeyAuthentication( |
| "snakanda", |
| IOUtils.toByteArray(HelixUtil.class.getClassLoader().getResourceAsStream("ssh/id_rsa")), |
| IOUtils.toByteArray(HelixUtil.class.getClassLoader().getResourceAsStream("ssh/id_rsa.pub")), |
| "dummy", |
| HelixUtil.class.getClassLoader().getResource("ssh/known_hosts").getPath(), |
| false |
| ); |
| |
| SSHServerInfo br2 = new SSHServerInfo("snakanda", "bigred2.uits.iu.edu", br2SshAuthentication,22); |
| |
| SSHTaskContext createDirTC = new SSHTaskContext(SSHTaskContext.TASK_TYPE.EXECUTE_COMMAND, |
| br2SshAuthentication, null, br2, "mkdir -p airavata"); |
| |
| SSHTaskContext qsubTC = new SSHTaskContext(SSHTaskContext.TASK_TYPE.EXECUTE_COMMAND, |
| br2SshAuthentication, null, br2, "qsub ~/airavata/job_tf.pbs"); |
| |
| |
| SSHTaskContext copyPbsTC = new SSHTaskContext(SSHTaskContext.TASK_TYPE.FILE_COPY, |
| br2SshAuthentication, null, br2, |
| HelixUtil.class.getClassLoader().getResource("job_tf.pbs").getPath(), "~/airavata/"); |
| |
| SSHTaskContext copyPyTC = new SSHTaskContext(SSHTaskContext.TASK_TYPE.FILE_COPY, |
| br2SshAuthentication, null, br2, |
| HelixUtil.class.getClassLoader().getResource("job_tf.pbs").getPath(), "~/airavata/"); |
| |
| ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, CREATE_DIR_TASK, getBytes(createDirTC)); |
| ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, COPY_PBS_TASK, getBytes(copyPbsTC)); |
| ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, COPY_PY_TASK, getBytes(copyPyTC)); |
| ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, QSUB_TASK, getBytes(qsubTC)); |
| } catch (Exception ex) { |
| ex.printStackTrace(); |
| return false; |
| } |
| return true; |
| } |
| |
| private static byte[] getBytes(Object object) throws Exception { |
| ByteArrayOutputStream bos = new ByteArrayOutputStream(); |
| ObjectOutput out = new ObjectOutputStream(bos); |
| out.writeObject(object); |
| out.flush(); |
| return bos.toByteArray(); |
| } |
| |
| public static Object getObject(byte[] objectBytes) throws Exception { |
| ByteArrayInputStream bis = new ByteArrayInputStream(objectBytes); |
| ObjectInput in = new ObjectInputStream(bis); |
| return in.readObject(); |
| } |
| |
| // public static void main(String[] args) throws Exception { |
| // |
| // CuratorFramework curatorClient = ZkUtils.getCuratorClient(); |
| // |
| // // set common data |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, USERNAME, "snakanda"); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, HOSTNAME, "bigred2.uits.iu.edu"); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, PORT, "22"); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, PUBLIC_KEY, |
| // IOUtils.toByteArray(HelixUtil.class.getClassLoader().getResourceAsStream("ssh/id_rsa.pub"))); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW, PRIVATE_KEY, |
| // IOUtils.toByteArray(HelixUtil.class.getClassLoader().getResourceAsStream("ssh/id_rsa"))); |
| // |
| // |
| // // set data for mkdir |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + CREATE_DIR_TASK, COMMAND, "mkdir -p airavata"); |
| // |
| // // set data for copy files |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + COPY_PBS_TASK, SRC_PATH, |
| // HelixUtil.class.getClassLoader().getResource("job_tf.pbs").getPath()); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + COPY_PBS_TASK, DEST_PATH, "~/airavata/"); |
| // |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + COPY_PY_TASK, SRC_PATH, |
| // HelixUtil.class.getClassLoader().getResource("code_tf.py").getPath()); |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + COPY_PY_TASK, DEST_PATH, "~/airavata/"); |
| // |
| // // set data for qsub |
| // ZkUtils.createZkNode(curatorClient, SSH_WORKFLOW + "/" + CREATE_DIR_TASK, COMMAND, "qsub ~/airavata/job_tf.pbs"); |
| // } |
| } |
