Update distribute_jobs.py

Added command-line option for number of CPUs and memory requested per job; transposed numpy array loaded from .npy array before initializing Data() object
olivercliff · Mar 5, 2022 · 5ba4904 · 5ba4904
1 parent 2b62be3
commit 5ba4904
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/distribute_jobs.py b/distribute_jobs.py
@@ -22,6 +22,12 @@
 parser.add_argument("--walltime_hrs", dest="walltime_hrs",
                     help = "OPTIONAL: Maximum walltime allowed for job. Default is 24 hours.",
                     default = "24")
+parser.add_argument("--cpu", dest="cpu",
+                    help = "OPTIONAL: Number of CPUs to request for each job. Default is 2.",
+                    default = "2")
+parser.add_argument("--mem", dest="mem",
+                    help = "OPTIONAL: Memory to request per job (in GB). Default is 8.",
+                    default = "8")
 parser.add_argument("--overwrite_pkl", dest="overwrite_pkl",
                     help = "OPTIONAL: overwrite all existing .pkl files in data directory? Default is False.",
                     default = False, action="store_true")
@@ -33,22 +39,24 @@
 user_email = args.user_email
 pbs_notify = args.pbs_notify
 walltime_hrs = args.walltime_hrs
+cpu = args.cpu
+mem = args.mem
 pyfile = os.path.abspath(args.compute_file)
 overwrite_pkl = args.overwrite_pkl
 
 # Import the rest of the modules
 from pyspi.calculator import Calculator
 from pyspi.data import Data
+import numpy as np
 import yaml
 import dill
-from shutil import copyfile
 from copy import deepcopy
 
 
 # Instantiate Calculator
 # Use user-generated config file if supplied to subset SPIs
 if args.pyspi_config is not None:
-    pyspi_config_file = args.pyspi_config
+    pyspi_config_file = os.path.abspath(args.pyspi_config)
     print(f"Custom config file: {pyspi_config_file}")
     basecalc = Calculator(configfile=pyspi_config_file)
 else:
@@ -71,7 +79,8 @@
                         name = str(config['name'])
                         labels = config['labels']
                         try:
-                            data = Data(data=file,dim_order=dim_order,name=name,normalise=True)
+                            data_from_file = np.load(file).transpose()
+                            data = Data(data=data_from_file,dim_order=dim_order,name=name,normalise=True)
                         except ValueError as err:
                             print(f'Issue loading dataset: {err}')
                             continue
@@ -81,7 +90,7 @@
                         calc.name = name
                         calc.labels = labels
                         sample_path = data_dir + "/" + name + "/"
-
+                        
                         # Create output directory
                         try:
                             os.mkdir(sample_path)
@@ -111,8 +120,8 @@
                             f.write("#!/bin/bash\n")
                             f.write(f"#PBS -N {name}\n")
                             f.write("#PBS -j oe\n")
-                            f.write(f"#PBS -o {data_dir}/{name}/pyspi_run.out\n")
-                            f.write("#PBS -l select=1:ncpus=2:mem=40GB\n")
+                            f.write(f"#PBS -o {data_dir}/{name}/pbsjob.out\n")
+                            f.write(f"#PBS -l select=1:ncpus={cpu}:mem={mem}GB\n")
                             f.write(f"#PBS -l walltime={walltime_hrs}:00:00\n")
                             f.write(f"#PBS -m {pbs_notify}\n")
                             if user_email is not None:
@@ -122,7 +131,7 @@
 
                             f.write("\n# --- CHANGE TO ANY RELEVANT CONDA INIT SCRIPTS\n")
                             f.write("module load Anaconda3-5.1.0\n")
-                            f.write("source /usr/physics/python/anaconda3/etc/profile.d/conda.csh\n")
+                            f.write("source /usr/physics/python/anaconda3/etc/profile.d/conda.sh\n")
                             f.write("# --- \n")
 
                             f.write("\nconda activate pyspi\n")
@@ -135,7 +144,7 @@
 
                             f.write("\n# Change to relevant directory and run our compute script\n")
                             f.write(f"cd {data_dir}\n")
-                            f.write(f"python {pyfile} {sample_pkl_output}")
+                            f.write(f"python {pyfile} {sample_pkl_output} > {data_dir}/{name}/pyspi_run.out")
 
 						# Submit the job
                         os.system(f'qsub {sample_pbs}')