bmtool 0.7.0.6.4__py3-none-any.whl → 0.7.1__py3-none-any.whl

bmtool/SLURM.py

@@ -1,12 +1,12 @@
-import time
-import os
-import subprocess
+import copy
 import json
-import requests
+import os
 import shutil
+import subprocess
 import time
-import copy
+
 import numpy as np
+import requests
 
 
 def check_job_status(job_id):
@@ -20,17 +20,17 @@ def check_job_status(job_id):
         str: The state of the job.
     """
     try:
-        result = subprocess.run(['scontrol', 'show', 'job', job_id], capture_output=True, text=True)
+        result = subprocess.run(["scontrol", "show", "job", job_id], capture_output=True, text=True)
         if result.returncode != 0:
             # this check is not needed if check_interval is less than 5 min (~300 seconds)
-            if 'slurm_load_jobs error: Invalid job id specified' in result.stderr:
-                return 'COMPLETED'  # Treat invalid job ID as completed because scontrol expires and removed job info when done.
-            #raise Exception(f"Error checking job status: {result.stderr}")
+            if "slurm_load_jobs error: Invalid job id specified" in result.stderr:
+                return "COMPLETED"  # Treat invalid job ID as completed because scontrol expires and removed job info when done.
+            # raise Exception(f"Error checking job status: {result.stderr}")
 
         job_state = None
-        for line in result.stdout.split('\n'):
-            if 'JobState=' in line:
-                job_state = line.strip().split('JobState=')[1].split()[0]
+        for line in result.stdout.split("\n"):
+            if "JobState=" in line:
+                job_state = line.strip().split("JobState=")[1].split()[0]
                 break
 
         if job_state is None:
@@ -39,7 +39,7 @@ def check_job_status(job_id):
         return job_state
     except Exception as e:
         print(f"Exception while checking job status: {e}", flush=True)
-        return 'UNKNOWN'
+        return "UNKNOWN"
 
 
 def submit_job(script_path):
@@ -55,29 +55,27 @@ def submit_job(script_path):
     Raises:
         Exception: If there is an error in submitting the job.
     """
-    result = subprocess.run(['sbatch', script_path], capture_output=True, text=True)
+    result = subprocess.run(["sbatch", script_path], capture_output=True, text=True)
     if result.returncode != 0:
         raise Exception(f"Error submitting job: {result.stderr}")
     job_id = result.stdout.strip().split()[-1]
     return job_id
 
 
-def send_teams_message(webhook,message):
+def send_teams_message(webhook, message):
     """Sends a message to a teams channel or chat
 
     Args:
         webhook (str): A microsoft teams webhook
         message (str): A message to send in the chat/channel
     """
-    message = {
-        "text": f"{message}" 
-    }
+    message = {"text": f"{message}"}
 
     # Send POST request to trigger the flow
     response = requests.post(
         webhook,
         json=message,  # Using 'json' instead of 'data' for automatic serialization
-        headers={'Content-Type': 'application/json'}
+        headers={"Content-Type": "application/json"},
     )
 
 
@@ -100,17 +98,20 @@ class seedSweep:
         Args:
             new_value: The new value for the parameter.
         """
-        with open(self.json_file_path, 'r') as f:
+        with open(self.json_file_path, "r") as f:
             data = json.load(f)
-        
+
         data[self.param_name] = new_value
-        
-        with open(self.json_file_path, 'w') as f:
+
+        with open(self.json_file_path, "w") as f:
             json.dump(data, f, indent=4)
-        
-        print(f"JSON file '{self.json_file_path}' modified successfully with {self.param_name}={new_value}.", flush=True)
-        
-    def change_json_file_path(self,new_json_file_path):
+
+        print(
+            f"JSON file '{self.json_file_path}' modified successfully with {self.param_name}={new_value}.",
+            flush=True,
+        )
+
+    def change_json_file_path(self, new_json_file_path):
         self.json_file_path = new_json_file_path
 
 
@@ -120,6 +121,7 @@ class multiSeedSweep(seedSweep):
     MultSeedSweeps are centered around some base JSON cell file. When that base JSON is updated, the other JSONs
     change according to their ratio with the base JSON.
     """
+
     def __init__(self, base_json_file_path, param_name, syn_dict, base_ratio=1):
         """
         Initializes the multipleSeedSweep instance.
@@ -144,22 +146,39 @@ class multiSeedSweep(seedSweep):
         self.edit_json(new_value)
         base_ratio = self.base_ratio
 
-        json_file_path = self.syn_dict_for_multi['json_file_path']
-        new_ratio = self.syn_dict_for_multi['ratio'] / base_ratio
-        
-        with open(json_file_path, 'r') as f:
+        json_file_path = self.syn_dict_for_multi["json_file_path"]
+        new_ratio = self.syn_dict_for_multi["ratio"] / base_ratio
+
+        with open(json_file_path, "r") as f:
             data = json.load(f)
         altered_value = new_ratio * new_value
         data[self.param_name] = altered_value
-    
-        with open(json_file_path, 'w') as f:
+
+        with open(json_file_path, "w") as f:
             json.dump(data, f, indent=4)
-    
-        print(f"JSON file '{json_file_path}' modified successfully with {self.param_name}={altered_value}.", flush=True)
+
+        print(
+            f"JSON file '{json_file_path}' modified successfully with {self.param_name}={altered_value}.",
+            flush=True,
+        )
 
 
 class SimulationBlock:
-    def __init__(self, block_name, time, partition, nodes, ntasks, mem, simulation_cases, output_base_dir,account=None,additional_commands=None,status_list = ['COMPLETED', 'FAILED', 'CANCELLED'],component_path=None):
+    def __init__(
+        self,
+        block_name,
+        time,
+        partition,
+        nodes,
+        ntasks,
+        mem,
+        simulation_cases,
+        output_base_dir,
+        account=None,
+        additional_commands=None,
+        status_list=["COMPLETED", "FAILED", "CANCELLED"],
+        component_path=None,
+    ):
         """
         Initializes the SimulationBlock instance.
 
@@ -172,10 +191,10 @@ class SimulationBlock:
             mem (int) : Number of gigabytes (per node)
             simulation_cases (dict): Dictionary of simulation cases with their commands.
             output_base_dir (str): Base directory for the output files.
-            account (str) : account to charge on HPC 
+            account (str) : account to charge on HPC
             additional commands (list): commands to run before bmtk model starts useful for loading modules
-            status_list (list): List of things to check before running next block. 
-                Adding RUNNING runs blocks faster but uses MUCH more resources and is only recommended on large HPC 
+            status_list (list): List of things to check before running next block.
+                Adding RUNNING runs blocks faster but uses MUCH more resources and is only recommended on large HPC
         """
         self.block_name = block_name
         self.time = time
@@ -185,7 +204,7 @@ class SimulationBlock:
         self.mem = mem
         self.simulation_cases = simulation_cases
         self.output_base_dir = output_base_dir
-        self.account = account 
+        self.account = account
         self.additional_commands = additional_commands if additional_commands is not None else []
         self.status_list = status_list
         self.job_ids = []
@@ -198,24 +217,33 @@ class SimulationBlock:
         Args:
             case_name (str): Name of the simulation case.
             command (str): Command to run the simulation.
-        
+
         Returns:
             str: Path to the batch script file.
         """
-        block_output_dir = os.path.join(self.output_base_dir, self.block_name)  # Create block-specific output folder
-        case_output_dir = os.path.join(block_output_dir, case_name)  # Create case-specific output folder
+        block_output_dir = os.path.join(
+            self.output_base_dir, self.block_name
+        )  # Create block-specific output folder
+        case_output_dir = os.path.join(
+            block_output_dir, case_name
+        )  # Create case-specific output folder
         os.makedirs(case_output_dir, exist_ok=True)
 
-        batch_script_path = os.path.join(block_output_dir, f'{case_name}_script.sh')
+        batch_script_path = os.path.join(block_output_dir, f"{case_name}_script.sh")
         additional_commands_str = "\n".join(self.additional_commands)
         # Conditional account linegit
         account_line = f"#SBATCH --account={self.account}\n" if self.account else ""
-        env_var_component_path = f"export COMPONENT_PATH={self.component_path}" if self.component_path else ""
-        mem_per_cpu = int(np.ceil(int(self.mem)/int(self.ntasks))) # do ceil cause more mem is always better then less
+        env_var_component_path = (
+            f"export COMPONENT_PATH={self.component_path}" if self.component_path else ""
+        )
+        mem_per_cpu = int(
+            np.ceil(int(self.mem) / int(self.ntasks))
+        )  # do ceil cause more mem is always better then less
 
         # Write the batch script to the file
-        with open(batch_script_path, 'w') as script_file:
-            script_file.write(f"""#!/bin/bash
+        with open(batch_script_path, "w") as script_file:
+            script_file.write(
+                f"""#!/bin/bash
 #SBATCH --job-name={self.block_name}_{case_name}
 #SBATCH --output={block_output_dir}/%x_%j.out
 #SBATCH --error={block_output_dir}/%x_%j.err
@@ -235,9 +263,10 @@ class SimulationBlock:
 export OUTPUT_DIR={case_output_dir}
 
 {command}
-""")
+"""
+            )
 
-        #print(f"Batch script created: {batch_script_path}", flush=True)
+        # print(f"Batch script created: {batch_script_path}", flush=True)
 
         return batch_script_path
 
@@ -247,7 +276,7 @@ export OUTPUT_DIR={case_output_dir}
         """
         for case_name, command in self.simulation_cases.items():
             script_path = self.create_batch_script(case_name, command)
-            result = subprocess.run(['sbatch', script_path], capture_output=True, text=True)
+            result = subprocess.run(["sbatch", script_path], capture_output=True, text=True)
             if result.returncode == 0:
                 job_id = result.stdout.strip().split()[-1]
                 self.job_ids.append(job_id)
@@ -267,7 +296,7 @@ export OUTPUT_DIR={case_output_dir}
             if status not in self.status_list:
                 return False
         return True
-    
+
     def check_block_completed(self):
         """checks if all the jobs in the block have been completed by slurm
 
@@ -276,8 +305,10 @@ export OUTPUT_DIR={case_output_dir}
         """
         for job_id in self.job_ids:
             status = check_job_status(job_id)
-            #print(f"status of job is {status}")
-            if status != 'COMPLETED': # can add PENDING here for debugging NOT FOR ACTUALLY USING IT 
+            # print(f"status of job is {status}")
+            if (
+                status != "COMPLETED"
+            ):  # can add PENDING here for debugging NOT FOR ACTUALLY USING IT
                 return False
         return True
 
@@ -289,10 +320,10 @@ export OUTPUT_DIR={case_output_dir}
         """
         for job_id in self.job_ids:
             status = check_job_status(job_id)
-            if status != 'RUNNING': # 
+            if status != "RUNNING":  #
                 return False
         return True
-    
+
     def check_block_submited(self):
         """checks if a job is running
 
@@ -301,7 +332,7 @@ export OUTPUT_DIR={case_output_dir}
         """
         for job_id in self.job_ids:
             status = check_job_status(job_id)
-            if status != 'PENDING': # 
+            if status != "PENDING":  #
                 return False
         return True
 
@@ -310,7 +341,9 @@ def get_relative_path(endpoint, absolute_path):
     """Convert absolute path to relative path for Globus transfer."""
     try:
         # Get the directories at the mount point
-        result = subprocess.run(["globus", "ls", f"{endpoint}:/"], capture_output=True, text=True, check=True)
+        result = subprocess.run(
+            ["globus", "ls", f"{endpoint}:/"], capture_output=True, text=True, check=True
+        )
         dirs = set(result.stdout.splitlines())  # Convert to a set for quicker lookup
 
         # Split the absolute path into parts
@@ -318,34 +351,34 @@ def get_relative_path(endpoint, absolute_path):
 
         # Find the first matching directory in the list
         for i, part in enumerate(path_parts):
-            if part+"/" in dirs:
+            if part + "/" in dirs:
                 # The mount point is everything up to and including this directory
                 mount_point = "/" + "/".join(path_parts[:i])
                 relative_path = absolute_path.replace(mount_point, "", 1).lstrip("/")
                 return relative_path
-        
+
         print("Error: Could not determine relative path.")
         return None
     except subprocess.CalledProcessError as e:
         print(f"Error retrieving directories from Globus: {e}")
         return None
 
+
 def globus_transfer(source_endpoint, dest_endpoint, source_path, dest_path):
     """
-    Transfers file using custom globus transfer function. 
+    Transfers file using custom globus transfer function.
     For more info see https://github.com/GregGlickert/transfer-files/blob/main/globus_transfer.sh
-    work in progress still... kinda forgot about this 
+    work in progress still... kinda forgot about this
     """
     relative_source_path = get_relative_path(source_endpoint, source_path)
     if relative_source_path is None:
         print("Transfer aborted: Could not determine relative source path.")
         return
-    
+
     command = f"globus transfer {source_endpoint}:{relative_source_path} {dest_endpoint}:{dest_path} --label 'bmtool slurm transfer'"
     os.system(command)
 
 
-
 class BlockRunner:
     """
     Class to handle submitting multiple blocks sequentially.
@@ -357,9 +390,17 @@ class BlockRunner:
         webhook (str): a microsoft webhook for teams. When used will send teams messages to the hook!
     """
 
-    def __init__(self, blocks, json_editor=None,json_file_path=None, param_name=None, 
-                 param_values=None, check_interval=60,syn_dict = None,
-                 webhook=None):
+    def __init__(
+        self,
+        blocks,
+        json_editor=None,
+        json_file_path=None,
+        param_name=None,
+        param_values=None,
+        check_interval=60,
+        syn_dict=None,
+        webhook=None,
+    ):
         self.blocks = blocks
         self.json_editor = json_editor
         self.param_values = param_values
@@ -377,12 +418,14 @@ class BlockRunner:
         for i, block in enumerate(self.blocks):
             print(block.output_base_dir)
             # Update JSON file with new parameter value
-            if self.json_file_path == None and self.param_values == None:
+            if self.json_file_path is None and self.param_values is None:
                 source_dir = block.component_path
                 destination_dir = f"{source_dir}{i+1}"
                 block.component_path = destination_dir
-                shutil.copytree(source_dir, destination_dir,dirs_exist_ok = True) # create new components folder 
-                print(f"skipping json editing for block {block.block_name}",flush=True)
+                shutil.copytree(
+                    source_dir, destination_dir, dirs_exist_ok=True
+                )  # create new components folder
+                print(f"skipping json editing for block {block.block_name}", flush=True)
             else:
                 if len(self.blocks) != len(self.param_values):
                     raise Exception("Number of blocks needs to each number of params given")
@@ -392,31 +435,34 @@ class BlockRunner:
                 destination_dir = f"{source_dir}{i+1}"
                 block.component_path = destination_dir
 
-                shutil.copytree(source_dir, destination_dir,dirs_exist_ok = True) # create new components folder
-                json_file_path = os.path.join(destination_dir,self.json_file_path)
-                
-                if self.syn_dict == None:
-                    json_editor = seedSweep(json_file_path , self.param_name)
+                shutil.copytree(
+                    source_dir, destination_dir, dirs_exist_ok=True
+                )  # create new components folder
+                json_file_path = os.path.join(destination_dir, self.json_file_path)
+
+                if self.syn_dict is None:
+                    json_editor = seedSweep(json_file_path, self.param_name)
                     json_editor.edit_json(new_value)
                 else:
                     # need to keep the orignal around
                     syn_dict_temp = copy.deepcopy(self.syn_dict)
-                    json_to_be_ratioed = syn_dict_temp['json_file_path']
-                    corrected_ratio_path = os.path.join(destination_dir,json_to_be_ratioed)
-                    syn_dict_temp['json_file_path'] = corrected_ratio_path
-                    json_editor = multiSeedSweep(json_file_path ,self.param_name,
-                                                syn_dict=syn_dict_temp,base_ratio=1)
-                    json_editor.edit_all_jsons(new_value) 
+                    json_to_be_ratioed = syn_dict_temp["json_file_path"]
+                    corrected_ratio_path = os.path.join(destination_dir, json_to_be_ratioed)
+                    syn_dict_temp["json_file_path"] = corrected_ratio_path
+                    json_editor = multiSeedSweep(
+                        json_file_path, self.param_name, syn_dict=syn_dict_temp, base_ratio=1
+                    )
+                    json_editor.edit_all_jsons(new_value)
 
             # Submit the block
             print(f"Submitting block: {block.block_name}", flush=True)
             block.submit_block()
             if self.webhook:
                 message = f"SIMULATION UPDATE: Block {i} has been submitted! There are {(len(self.blocks)-1)-i} left to be submitted"
-                send_teams_message(self.webhook,message)
+                send_teams_message(self.webhook, message)
 
             # Wait for the block to complete
-            if i == len(self.blocks) - 1:  
+            if i == len(self.blocks) - 1:
                 while not block.check_block_status():
                     print(f"Waiting for the last block {i} to complete...")
                     time.sleep(self.check_interval)
@@ -424,12 +470,12 @@ class BlockRunner:
                 while not block.check_block_status():
                     print(f"Waiting for block {i} to complete...")
                     time.sleep(self.check_interval)
-            
+
             print(f"Block {block.block_name} completed.", flush=True)
-        print("All blocks are done!",flush=True)
+        print("All blocks are done!", flush=True)
         if self.webhook:
             message = "SIMULATION UPDATE: Simulation are Done!"
-            send_teams_message(self.webhook,message)
+            send_teams_message(self.webhook, message)
 
     def submit_blocks_parallel(self):
         """
@@ -437,47 +483,54 @@ class BlockRunner:
         Also the json_file_path should be the path after the components dir
         """
         for i, block in enumerate(self.blocks):
-            if self.param_values == None:
+            if self.param_values is None:
                 source_dir = block.component_path
                 destination_dir = f"{source_dir}{i+1}"
                 block.component_path = destination_dir
-                shutil.copytree(source_dir, destination_dir,dirs_exist_ok = True) # create new components folder 
-                print(f"skipping json editing for block {block.block_name}",flush=True)
+                shutil.copytree(
+                    source_dir, destination_dir, dirs_exist_ok=True
+                )  # create new components folder
+                print(f"skipping json editing for block {block.block_name}", flush=True)
             else:
-                if block.component_path == None:
-                    raise Exception("Unable to use parallel submitter without defining the component path")
+                if block.component_path is None:
+                    raise Exception(
+                        "Unable to use parallel submitter without defining the component path"
+                    )
                 new_value = self.param_values[i]
-                
+
                 source_dir = block.component_path
                 destination_dir = f"{source_dir}{i+1}"
                 block.component_path = destination_dir
 
-                shutil.copytree(source_dir, destination_dir,dirs_exist_ok = True) # create new components folder 
-                json_file_path = os.path.join(destination_dir,self.json_file_path)
-                
-                if self.syn_dict == None:
-                    json_editor = seedSweep(json_file_path , self.param_name)
+                shutil.copytree(
+                    source_dir, destination_dir, dirs_exist_ok=True
+                )  # create new components folder
+                json_file_path = os.path.join(destination_dir, self.json_file_path)
+
+                if self.syn_dict is None:
+                    json_editor = seedSweep(json_file_path, self.param_name)
                     json_editor.edit_json(new_value)
                 else:
                     # need to keep the orignal around
                     syn_dict_temp = copy.deepcopy(self.syn_dict)
-                    json_to_be_ratioed = syn_dict_temp['json_file_path']
-                    corrected_ratio_path = os.path.join(destination_dir,json_to_be_ratioed)
-                    syn_dict_temp['json_file_path'] = corrected_ratio_path
-                    json_editor = multiSeedSweep(json_file_path ,self.param_name,
-                                                syn_dict_temp,base_ratio=1)
-                    json_editor.edit_all_jsons(new_value) 
-                # submit block with new component path 
+                    json_to_be_ratioed = syn_dict_temp["json_file_path"]
+                    corrected_ratio_path = os.path.join(destination_dir, json_to_be_ratioed)
+                    syn_dict_temp["json_file_path"] = corrected_ratio_path
+                    json_editor = multiSeedSweep(
+                        json_file_path, self.param_name, syn_dict_temp, base_ratio=1
+                    )
+                    json_editor.edit_all_jsons(new_value)
+                # submit block with new component path
             print(f"Submitting block: {block.block_name}", flush=True)
             block.submit_block()
             if i == len(self.blocks) - 1:
-                print("\nEverything has been submitted. You can close out of this or keep this script running to get a message when everything is finished\n")
+                print(
+                    "\nEverything has been submitted. You can close out of this or keep this script running to get a message when everything is finished\n"
+                )
                 while not block.check_block_status():
                     print(f"Waiting for the last block {i} to complete...")
                     time.sleep(self.check_interval)
-                
+
         if self.webhook:
             message = "SIMULATION UPDATE: Simulations are Done!"
-            send_teams_message(self.webhook,message)
-            
-            
+            send_teams_message(self.webhook, message)

bmtool/__init__.py

@@ -3,4 +3,4 @@ __version__ = "0.0.1"
 import warnings
 
 with warnings.catch_warnings():
-        warnings.filterwarnings("ignore",category=FutureWarning)
+    warnings.filterwarnings("ignore", category=FutureWarning)

bmtool/__main__.py

@@ -1,27 +1,28 @@
 import argparse
-from argparse import RawTextHelpFormatter,SUPPRESS
 import sys
+from argparse import SUPPRESS, RawTextHelpFormatter
 
 use_description = """
 Build, plot or debug BMTK models easily.
 
 python -m bmtool.build
-python -m bmtool.plot 
+python -m bmtool.plot
 python -m bmtool.debug
 python -m bmtool.util
 
 """
 
-if __name__ == '__main__':
-    
-    parser = argparse.ArgumentParser(description=use_description, formatter_class=RawTextHelpFormatter,usage=SUPPRESS)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description=use_description, formatter_class=RawTextHelpFormatter, usage=SUPPRESS
+    )
     options = None
     try:
         if not len(sys.argv) > 1:
             raise
-        if sys.argv[1] in ['-h','--h','-help','--help','help']:
+        if sys.argv[1] in ["-h", "--h", "-help", "--help", "help"]:
             raise
         options = parser.parse_args()
     except:
         parser.print_help()
-        sys.exit(0)
+        sys.exit(0)