From 22c11b54b96780606eebf7ec31f796c263a74046 Mon Sep 17 00:00:00 2001 From: Justin Hiemstra Date: Tue, 17 Sep 2024 20:53:39 +0000 Subject: [PATCH] Incorporate review feedback --- docker-wrappers/SPRAS/snakemake_long.py | 46 ++++++++++++++----------- 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/docker-wrappers/SPRAS/snakemake_long.py b/docker-wrappers/SPRAS/snakemake_long.py index 787f30c8..1f851409 100755 --- a/docker-wrappers/SPRAS/snakemake_long.py +++ b/docker-wrappers/SPRAS/snakemake_long.py @@ -20,28 +20,33 @@ need different arguments. The "top" mode is for submitting the HTCondor wrapper, and the "long" mode is for running the Snakemake command itself. """ -def parseArgs(isLocal=False): +def parse_args(isLocal=False): parser = argparse.ArgumentParser(description="A tool for long-running Snakemake jobs with HTCondor.") - # We add a special command to trigger allowing this script to execute the long-running Snakemake command. if isLocal: + # We add a special command that allows this singular executable to serve two purposes. The executable + # is first run by the user with their args to submit the local universe job. Then, the local universe + # job runs `snakemake_long.py long ` to indicate to the script that it's time to submit the + # long-running Snakemake process instead of submitting another local universe job. parser.add_argument("command", help="Helper command to run", choices=["long"]) - parser.add_argument("--snakefile", help="The Snakefile to run.", required=False) + parser.add_argument("--snakefile", help="The Snakefile to run. If omitted, the Snakefile is assumed to be in the current directory.", required=False) parser.add_argument("--profile", help="A path to a directory containing the desired Snakemake profile.", required=True) # I'd love to change this to "logdir", but using the same name as Snakemake for consistency of feeling between this script # and Snakemake proper. - parser.add_argument("--htcondor-jobdir", help="The directory Snakemake will write logs to.", required=False) + parser.add_argument("--htcondor-jobdir", help="The directory Snakemake will write logs to. If omitted, a 'logs` directory will be created in the current directory", required=False) return parser.parse_args() """ Given a Snakefile, profile, and HTCondor job directory, submit a local universe job that runs Snakemake from the context of the submission directory. """ -def submitLocal(snakefile, profile, htcondor_jobdir): +def submit_local(snakefile, profile, htcondor_jobdir): # Get the location of this script, which also serves as the executable for the condor job. script_location = pathlib.Path(__file__).resolve() submit_description = htcondor.Submit({ "executable": script_location, + # We use the "long" command to indicate to the script that it should run the Snakemake command instead of submitting another job. + # See comment in parse_args for more information. "arguments": f"long --snakefile {snakefile} --profile {profile} --htcondor-jobdir {htcondor_jobdir}", "universe": "local", "request_disk": "512MB", @@ -49,16 +54,16 @@ def submitLocal(snakefile, profile, htcondor_jobdir): "request_memory": 512, # Set up logging - "log": f"{htcondor_jobdir}/snakemake-long.log", - "output": f"{htcondor_jobdir}/snakemake-long.out", - "error": f"{htcondor_jobdir}/snakemake-long.err", + "log": f"{htcondor_jobdir}/snakemake.log", + "output": f"{htcondor_jobdir}/snakemake.out", + "error": f"{htcondor_jobdir}/snakemake.err", # Specify `getenv` so that our script uses the appropriate environment # when it runs in local universe. This allows the job to access # modules we've installed in the submission environment (notably spras). "getenv": "true", - "JobBatchName": f"spras-long-{time.strftime('%Y%m%d-%H%M%S')}", + "JobBatchName": f"spras-{time.strftime('%Y%m%d-%H%M%S')}", }) schedd = htcondor.Schedd() @@ -70,8 +75,8 @@ def submitLocal(snakefile, profile, htcondor_jobdir): The top level function for the script that handles file creation/validation and triggers submission of the wrapper job. """ -def topMain(): - args = parseArgs() +def top_main(): + args = parse_args() # Check if the snakefile is provided. If not, assume it's in the current directory. if args.snakefile is None: @@ -97,19 +102,18 @@ def topMain(): os.makedirs(args.htcondor_jobdir) - submitLocal(args.snakefile, args.profile, args.htcondor_jobdir) + submit_local(args.snakefile, args.profile, args.htcondor_jobdir) return 0 """ - +Command to activate conda environment and run Snakemake. This is run by the local universe job, not the user. """ -def longMain(): - args = parseArgs(True) +def long_main(): + args = parse_args(True) - # Command to activate conda environment and run Snakemake. Note that we need to unset APPTAINER_CACHEDIR - # in this case but not in the local terminal case because the wrapper HTCondor job has a different environment - # and populating this value causes Snakemake to fail when it tries to write to spool (a read-only filesystem from - # the perspective of the EP job). + # Note that we need to unset APPTAINER_CACHEDIR in this case but not in the local terminal case because the wrapper + # HTCondor job has a different environment and populating this value causes Snakemake to fail when it tries to write + # to spool (a read-only filesystem from the perspective of the EP job). command = f""" source $(conda info --base)/etc/profile.d/conda.sh && \ conda activate spras && \ @@ -126,9 +130,9 @@ def longMain(): def main(): if len(sys.argv) > 1: if sys.argv[1] in ["long"]: - return longMain() + return long_main() - return topMain() + return top_main() if __name__ == '__main__': sys.exit(main())