Fix spaces handling in array job arguments (#45)

* Use shell arrays for command arguments in submission script * Add integration test for arguments containing spaces Signed-off-by: yichen88 <[email protected]>
powsybl · May 15, 2020 · 10b8c0e · 10b8c0e
1 parent a19a74f
commit 10b8c0e
Show file tree

Hide file tree

Showing 9 changed files with 82 additions and 20 deletions.
diff --git a/...tion-slurm-test/src/test/java/com/powsybl/computation/slurm/SlurmNormalExecutionTest.java b/...tion-slurm-test/src/test/java/com/powsybl/computation/slurm/SlurmNormalExecutionTest.java
@@ -131,6 +131,27 @@ public List<CommandExecution> before(Path workingDir) {
         baseTest(supplier);
     }
 
+    @Test
+    public void testArgsWithSpace() {
+        Supplier<AbstractExecutionHandler<String>> supplier = () -> new AbstractReturnOKExecutionHandler() {
+            @Override
+            public List<CommandExecution> before(Path workingDir) {
+                return CommandExecutionsTestFactory.argsWithSpaces(3);
+            }
+
+            @Override
+            public String after(Path workingDir, ExecutionReport report) throws IOException {
+                if (Files.exists(workingDir.resolve("line 1,line 2")) && Files.exists(workingDir.resolve("v2"))) {
+                    return super.after(workingDir, report);
+                } else {
+                    failed = true;
+                    return "KO";
+                }
+            }
+        };
+        baseTest(supplier);
+    }
+
     @Test
     public void testTwoSimpleCmd() {
         Supplier<AbstractExecutionHandler<String>> supplier = () -> new AbstractReturnOKExecutionHandler() {

diff --git a/computation-slurm/src/main/java/com/powsybl/computation/slurm/CommandUtils.java b/computation-slurm/src/main/java/com/powsybl/computation/slurm/CommandUtils.java
@@ -6,7 +6,7 @@
  */
 package com.powsybl.computation.slurm;
 
-import java.util.Collection;
+import java.util.List;
 import java.util.StringJoiner;
 import java.util.stream.Collector;
 
@@ -31,7 +31,7 @@ private static Collector<String, StringJoiner, String> getWrapperAndJoiner() {
     /**
      * Generates a command string, with each argument wrapped with quotes.
      */
-    static String commandToString(String program, Collection<String> args) {
+    static String commandToString(String program, List<String> args) {
         requireNonNull(program);
         requireNonNull(args);
 
@@ -44,7 +44,7 @@ static String commandToString(String program, Collection<String> args) {
      * @param args
      * @return the argu's string
      */
-    static String commandArgsToString(Collection<String> args) {
+    static String commandArgsToString(List<String> args) {
         requireNonNull(args);
         return args.stream().collect(getWrapperAndJoiner());
     }

diff --git a/computation-slurm/src/main/java/com/powsybl/computation/slurm/SbatchScriptGenerator.java b/computation-slurm/src/main/java/com/powsybl/computation/slurm/SbatchScriptGenerator.java
@@ -37,6 +37,7 @@ class SbatchScriptGenerator {
     private static final String INDENTATION_6 = "      ";
     private static final String SH_CASE_BREAK = INDENTATION_4 + ";;";
     private static final String SPL_CMD_ARGS = "ARGS";
+    private static final String CALL_SPL_CMD_ARGS = " \"${ARGS[@]}\"";
     private static final String SUB_CMD_ARGS = "ARGS_";
     private static final String PRE_FILE = "PRE";
     private static final String POST_FILE = "POST";
@@ -213,7 +214,7 @@ private void arrayJobCase(List<String> shell, CommandExecution commandExecution)
                 case SIMPLE:
                     SimpleCommand simpleCmd = (SimpleCommand) command;
                     String args = CommandUtils.commandArgsToString(simpleCmd.getArgs(caseIdx));
-                    shell.add(INDENTATION_6 + SPL_CMD_ARGS + "=\"" + args + "\"");
+                    shell.add(INDENTATION_6 + SPL_CMD_ARGS + "=(" + args + ")");
                     shell.add(SH_CASE_BREAK);
                     break;
                 case GROUP:
@@ -223,7 +224,7 @@ private void arrayJobCase(List<String> shell, CommandExecution commandExecution)
                     for (int i = 0; i < subCommands.size(); i++) {
                         GroupCommand.SubCommand cmd = subCommands.get(i);
                         String argsSub = CommandUtils.commandArgsToString(cmd.getArgs(caseIdx));
-                        String de = SUB_CMD_ARGS + i + "=\"" + argsSub + "\"";
+                        String de = SUB_CMD_ARGS + i + "=(" + argsSub + ")";
                         subArgs.add(de);
                     }
                     String subArgsJoined = String.join(" ", subArgs);
@@ -293,11 +294,12 @@ private void cmd(List<String> shell, Command command, Path workingDir) {
     }
 
     private void simpleCmdWithArgs(List<String> list, SimpleCommand simpleCommand) {
-        list.add(simpleCommand.getProgram() + " $" + SPL_CMD_ARGS);
+        list.add(simpleCommand.getProgram() + CALL_SPL_CMD_ARGS);
     }
 
     private void subCmdWithArgs(List<String> list, GroupCommand.SubCommand subCommand, int idxInGroup) {
-        list.add(subCommand.getProgram() + " $" + SUB_CMD_ARGS + idxInGroup);
+        // cmd "${ARGS_1{[@]}"
+        list.add(subCommand.getProgram() + " \"${ARGS_" + idxInGroup + "[@]}\"");
     }
 
 }
diff --git a/...ation-slurm/src/test/java/com/powsybl/computation/slurm/CommandExecutionsTestFactory.java b/...ation-slurm/src/test/java/com/powsybl/computation/slurm/CommandExecutionsTestFactory.java
@@ -41,6 +41,17 @@ static List<CommandExecution> simpleCmdWithCount(int executionCount) {
         return Collections.singletonList(commandExecution);
     }
 
+    static List<CommandExecution> argsWithSpaces(int executionCount) {
+        Command command = new SimpleCommandBuilder()
+                .id("spaces")
+                .program("touch")
+                .args(i -> Arrays.asList("line 1,line 2", "v2"))
+                .timeout(60)
+                .build();
+        CommandExecution commandExecution = new CommandExecution(command, executionCount);
+        return Collections.singletonList(commandExecution);
+    }
+
     /**
      * The 4th batch would fail.
      * @param executionCount

diff --git a/computation-slurm/src/test/java/com/powsybl/computation/slurm/SbatchScriptGeneratorTest.java b/computation-slurm/src/test/java/com/powsybl/computation/slurm/SbatchScriptGeneratorTest.java
@@ -79,6 +79,20 @@ public void testSimpleCmd() {
                 "touch /tmp/flags/mydone_workingPath_12345_$SLURM_JOBID"), shell);
     }
 
+    @Test
+    public void testArgsWithSpaces() {
+        List<CommandExecution> commandExecutions = CommandExecutionsTestFactory.argsWithSpaces(3);
+        CommandExecution commandExecution = commandExecutions.get(commandIdx);
+        Command command = commandExecution.getCommand();
+        List<String> shell = new SbatchScriptGenerator(flagPath).parser(command, 0, workingPath, Collections.emptyMap());
+        // not array job
+        assertEquals(ImmutableList.of("#!/bin/sh",
+                "touch \"line 1,line 2\" \"v2\"",
+                "rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_JOBID; exit $rc; fi",
+                "touch /tmp/flags/mydone_workingPath_12345_$SLURM_JOBID"), shell);
+        assertCommandExecutionToShell(commandExecutions.get(0), "argsWithSpaces.batch");
+    }
+
     @Test
     public void testSimpleCmdWithCount() {
         List<CommandExecution> commandExecutions = CommandExecutionsTestFactory.simpleCmdWithCount(3);

diff --git a/computation-slurm/src/test/resources/expectedShell/argsWithSpaces.batch b/computation-slurm/src/test/resources/expectedShell/argsWithSpaces.batch
@@ -0,0 +1,14 @@
+#!/bin/sh
+case $SLURM_ARRAY_TASK_ID in
+    0)
+      ARGS=("line 1,line 2" "v2")
+    ;;
+    1)
+      ARGS=("line 1,line 2" "v2")
+    ;;
+    2)
+      ARGS=("line 1,line 2" "v2")
+    ;;
+esac
+touch "${ARGS[@]}"
+rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID; exit $rc; fi
diff --git a/computation-slurm/src/test/resources/expectedShell/group3.batch b/computation-slurm/src/test/resources/expectedShell/group3.batch
@@ -1,16 +1,16 @@
 #!/bin/sh
 case $SLURM_ARRAY_TASK_ID in
     0)
-      ARGS_0=""5s"" ARGS_1=""sub2""
+      ARGS_0=("5s") ARGS_1=("sub2")
     ;;
     1)
-      ARGS_0=""5s"" ARGS_1=""sub2""
+      ARGS_0=("5s") ARGS_1=("sub2")
     ;;
     2)
-      ARGS_0=""5s"" ARGS_1=""sub2""
+      ARGS_0=("5s") ARGS_1=("sub2")
     ;;
 esac
-sleep $ARGS_0
+sleep "${ARGS_0[@]}"
 rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID; exit $rc; fi
-echo $ARGS_1
+echo "${ARGS_1[@]}"
 rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID; exit $rc; fi
diff --git a/computation-slurm/src/test/resources/expectedShell/myEchoSimpleCmdWithUnzipZip.batch b/computation-slurm/src/test/resources/expectedShell/myEchoSimpleCmdWithUnzipZip.batch
@@ -3,20 +3,20 @@ case $SLURM_ARRAY_TASK_ID in
     0)
       PRE0=in0.zip
       POST0=out0
-      ARGS=""in0" "out0""
+      ARGS=("in0" "out0")
     ;;
     1)
       PRE0=in1.zip
       POST0=out1
-      ARGS=""in1" "out1""
+      ARGS=("in1" "out1")
     ;;
     2)
       PRE0=in2.zip
       POST0=out2
-      ARGS=""in2" "out2""
+      ARGS=("in2" "out2")
     ;;
 esac
 unzip -o -q $PRE0
-/home/dev-itesla/myapps/myecho.sh $ARGS
+/home/dev-itesla/myapps/myecho.sh "${ARGS[@]}"
 rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID; exit $rc; fi
 gzip $POST0
diff --git a/computation-slurm/src/test/resources/expectedShell/simpleCmdWithCount3.batch b/computation-slurm/src/test/resources/expectedShell/simpleCmdWithCount3.batch
@@ -1,14 +1,14 @@
 #!/bin/sh
 case $SLURM_ARRAY_TASK_ID in
     0)
-      ARGS=""te1st0""
+      ARGS=("te1st0")
     ;;
     1)
-      ARGS=""te1st1""
+      ARGS=("te1st1")
     ;;
     2)
-      ARGS=""te1st2""
+      ARGS=("te1st2")
     ;;
 esac
-echo $ARGS
+echo "${ARGS[@]}"
 rc=$?; if [[ $rc != 0 ]]; then touch /tmp/flags/myerror_workingPath_12345_$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID; exit $rc; fi