From f986551dca5bc0a182ccf79e041507143cc095e3 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Mon, 9 Sep 2024 15:09:44 -0500
Subject: [PATCH] Use `shlex.quote` to convert the arguments into unix format

---
 alea/submitters/htcondor.py          | 23 +++++++++++------------
 alea/submitters/run_toymc_wrapper.sh |  6 ++++++
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/alea/submitters/htcondor.py b/alea/submitters/htcondor.py
index 0efb18b..abc2a0c 100644
--- a/alea/submitters/htcondor.py
+++ b/alea/submitters/htcondor.py
@@ -2,7 +2,6 @@
 import getpass
 import tarfile
 import shlex
-import json
 import tempfile
 import time
 import threading
@@ -25,6 +24,7 @@
     TransformationCatalog,
     ReplicaCatalog,
 )
+from alea.runner import Runner
 from alea.submitter import Submitter
 from alea.utils import load_yaml, dump_yaml
 
@@ -449,13 +449,14 @@ def _initialize_job(
 
         """
         job = Job(name)
-        job.add_profiles(Namespace.CONDOR, "request_cpus", f"{cores}")
 
         if run_on_submit_node:
             job.add_selector_profile(execution_site="local")
             # no other attributes on a local job
             return job
 
+        job.add_profiles(Namespace.CONDOR, "request_cpus", f"{cores}")
+
         # Set memory and disk requirements
         # If the job fails, retry with more memory and disk
         # Somehow we need to write memory in MB and disk in kB
@@ -497,13 +498,7 @@ def _add_separate_job(self, combine_i):
         """Add a separate job to the workflow."""
         logger.info(f"Adding separate job {combine_i} to the workflow")
         separate_name = "separate"
-        separate_job = self._initialize_job(
-            name=separate_name,
-            cores=1,
-            memory=self.request_memory * 2,
-            disk=self.combine_disk,
-            run_on_submit_node=True,
-        )
+        separate_job = self._initialize_job(name=separate_name, run_on_submit_node=True)
 
         # Separate job configuration: all toymc results and files will be combined into one tarball
         separate_job.add_inputs(File(f"{self.workflow_id}-{combine_i}-combined_output.tar.gz"))
@@ -565,6 +560,9 @@ def _generate_workflow(self, name="run_toymc_wrapper"):
         combine_i = 0
         new_to_combine = True
 
+        # Prepare for argument conversion
+        _, _, annotations = Runner.runner_arguments()
+
         # Generate jobstring and output names from tickets generator
         for job_id, (script, _) in enumerate(self.combined_tickets_generator()):
             # If the number of jobs to combine is reached, add a new combine job
@@ -638,10 +636,11 @@ def _generate_workflow(self, name="run_toymc_wrapper"):
 
             # Add the arguments into the job
             # Using escaped argument to avoid the shell syntax error
-            def _extract_all_to_tuple(d):
+            def _extract_all_to_tuple(kwargs) -> tuple:
+                """Generate the submission script from the runner arguments."""
                 return tuple(
-                    f"{json.dumps(str(d[key])).replace(' ', '')}".replace("'", '\\"')
-                    for key in d.keys()
+                    shlex.quote(Submitter.arg_to_str(kwargs[arg], annotation))
+                    for arg, annotation in annotations.items()
                 )
 
             # Correct the paths in the arguments
diff --git a/alea/submitters/run_toymc_wrapper.sh b/alea/submitters/run_toymc_wrapper.sh
index c7ed324..4bebc4b 100644
--- a/alea/submitters/run_toymc_wrapper.sh
+++ b/alea/submitters/run_toymc_wrapper.sh
@@ -2,6 +2,12 @@
 
 set -e
 
+# Check if number of arguments passed is correct
+if [ $# -ne 21 ]; then
+    echo "Error: You need to provide required number of arguments."
+    exit 1
+fi
+
 # Extract the arguments
 statistical_model=$1
 poi=$2