Merge pull request #48 from brockpalen/new/wrapper

New/wrapper

bin/.archivetar

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+# Brock Palen
+# brockp@umich.edu
+# 7/2020
+#
+#  prep a directory for placement in dataden
+#  process:
+#    1. run mpiFileUtils / dwalk  (deafault sort in name / path order) all files < minsize
+#    2. Take resulting list build tar lists by summing size until > tarsize (before compression)
+#    3. Tar each list:  OR --dryrun create list with est size
+#       a. Create Index file of contents
+#       b. Optionally compress -z / -j  with gzip/pigz bzip/lbzip2 if installed
+#       c. Optionally purge
+#    4. (?) Kick out optimized untar script (pigz / lbzip2)
+
+import sys
+
+import archivetar
+
+archivetar.main(sys.argv)

bin/archivetar

@@ -1,21 +1,52 @@
-#!/usr/bin/env python3
-
-# Brock Palen
-# brockp@umich.edu
-# 7/2020
-#
-#  prep a directory for placement in dataden
-#  process:
-#    1. run mpiFileUtils / dwalk  (deafault sort in name / path order) all files < minsize
-#    2. Take resulting list build tar lists by summing size until > tarsize (before compression)
-#    3. Tar each list:  OR --dryrun create list with est size
-#       a. Create Index file of contents
-#       b. Optionally compress -z / -j  with gzip/pigz bzip/lbzip2 if installed
-#       c. Optionally purge
-#    4. (?) Kick out optimized untar script (pigz / lbzip2)
-
+import os
+import subprocess
 import sys
+import shutil
+
+def main():
+    '''
+    To disable slurm execution entirely, unset env-variable AT_SLURM_OFFLOAD. 
+    To enable slurm execution, set env-variable AT_SLURM_OFFLOAD=1
+
+    ARCHIVETAR_TASKS, ARCHIVETAR_MEM, and ARCHIVETAR_PAR control the cores, memory and partition requirements
+        of the SLURM job srun executes.
+    '''
+    # Check for help options and run locally
+    if len(sys.argv) > 1 and (sys.argv[1] == "-h" or sys.argv[1] == "--help"):
+        print("\033[34m==>\033[32m Running locally for help option\033[0m")
+        subprocess.run([".archivetar", "-h"])
+        sys.exit(0)
+
+    # Handle the case where no arguments are provided (just "archivetar")
+    if len(sys.argv) == 1:
+        print("\033[34m==>\033[32m Running archivetar with no arguments\033[0m")
+        subprocess.run([".archivetar"])
+        sys.exit(0)
+
+    # Check if running inside a SLURM job
+    slurm_job_id = os.getenv("SLURM_JOB_ID")
+    at_slurm_offload = os.getenv("AT_SLURM_OFFLOAD")
+
+    if slurm_job_id:
+        # Run locally inside SLURM without executing another SLURM job.
+        result = subprocess.run([".archivetar"] + sys.argv[1:])
+        sys.exit(result.returncode)
+    elif at_slurm_offload and shutil.which("srun"):
+            print("\033[34m==>\033[35m Running archivetar within SLURM\033[0m")
+            # Get environment variables or use default values
+            tasks = os.getenv("AT_TASKS", "8")  # Default to 8 tasks if ARCHIVETAR_TASKS is not set
+            mem = os.getenv("AT_MEM", "40G")    # Default to 40G if ARCHIVETAR_MEM is not set
+            partition = os.getenv("AT_PAR", "archive") # Default to archive if ARCHIVETAR_PAR is not set
+
+            # Run Python script from within SLURM
+            cmd = f"srun --partition={partition} --cpu-bind=no --ntasks=1 --cpus-per-task={tasks} --mem={mem} --job-name=archivetar_{os.getenv('USER')} --time=14-00:00:00 --pty bash -c '.archivetar {' '.join(sys.argv[1:])}'"
+            result = subprocess.run(cmd, shell=True)
+            sys.exit(result.returncode)
+    else:
+        # Run locally without SLURM
+        result = subprocess.run([".archivetar"] + sys.argv[1:])
+        sys.exit(result.returncode)
 
-import archivetar
 
-archivetar.main(sys.argv)
+if __name__ == "__main__":
+    main()