| import re
|
| from functools import lru_cache
|
| import subprocess
|
| from typing import List, Optional, Dict, Tuple
|
| import os
|
|
|
| from libs.utils_classes import AlphaFoldJobInfo, RunAlphaFoldResult
|
|
|
| |
| MAX_AF_SIZE = 3000 |
| REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) |
| ASSEMBLER_DIR = os.path.join(REPO_ROOT, "model", "CombinatorialAssembler") |
| AF2TRANS_BIN_PATH = os.path.join(ASSEMBLER_DIR, "AF2trans.out") |
| COMB_ASSEMBLY_BIN_PATH = os.path.join(ASSEMBLER_DIR, "CombinatorialAssembler.out") |
|
|
|
|
|
|
| @lru_cache(1)
|
| def get_squeue_output():
|
| return subprocess.check_output(["squeue"])
|
|
|
|
|
| def get_slurm_ids_from_log(slurms_log: str) -> List[str]:
|
| if not os.path.exists(slurms_log):
|
| return []
|
| return re.findall(r"Submitted batch job (\d+)", open(slurms_log, "r").read())
|
|
|
|
|
| def check_for_active_slurm(slurms_log: str) -> bool:
|
| slurm_ids = get_slurm_ids_from_log(slurms_log)
|
| if len(slurm_ids) == 0:
|
| return False
|
| squeue = get_squeue_output()
|
| for slurm_id in slurm_ids:
|
| for line in squeue.split(b"\n"):
|
| if slurm_id.encode() in line:
|
| return True
|
| return False
|
|
|
|
|
| def get_last_slurm_output(output_folder: str, slurms_log: str) -> Optional[str]:
|
| slurm_ids = get_slurm_ids_from_log(slurms_log)
|
| if len(slurm_ids) == 0:
|
| return None
|
| if not os.path.exists(os.path.join(output_folder, f"slurm-{slurm_ids[-1]}.out")):
|
| return None
|
| return open(os.path.join(output_folder, f"slurm-{slurm_ids[-1]}.out"), "r").read()
|
|
|
|
|
| def is_bash_running(bash_path: str):
|
| output_path = os.path.dirname(bash_path)
|
| slurm_ids_log_path = os.path.join(output_path, f"slurm_logs_{os.path.basename(bash_path)}.log")
|
| if check_for_active_slurm(slurm_ids_log_path):
|
| return True
|
| return False
|
|
|
|
|
| def run_bash_file(bash_path: str):
|
| output_path = os.path.dirname(bash_path)
|
| slurm_ids_log_path = os.path.join(output_path, f"slurm_logs_{os.path.basename(bash_path)}.log")
|
| subprocess.run(f"sbatch -c2 --mem=10G --time=3:0:0 {bash_path} >> {slurm_ids_log_path}", shell=True)
|
|
|
|
|
|
|
| def get_scores_from_path(af_pdb_path: str):
|
| return af_pdb_path.replace("unrelaxed", "scores").replace("pdb", "json")
|
|
|
|
|
| def run_alphafold_on_fasta(fasta_path: str, output_folder: str) -> Tuple[RunAlphaFoldResult, List[str]]:
|
| """
|
| This function gets a path to a fasta and an output folder and runs alphafold on the fasta.
|
| The function return a tuple of the status of the alphafold job(running/completed/failed) and if completed,
|
| the second part of the returned tuple should be a list of paths to the alphafold predicted pdbs.
|
| """
|
| af_jobname = [i for i in open(fasta_path, "r").read().split("\n") if i.startswith(">")][0][1:]
|
|
|
| os.chdir(output_folder)
|
| cluster_run_path = os.path.join(output_folder, f"cluster_run_{af_jobname}.sh")
|
| slurm_ids_log_path = os.path.join(output_folder, f"slurm_logs_{af_jobname}.sh")
|
|
|
|
|
| result_files = [os.path.join(output_folder, filename) for filename in os.listdir(output_folder)
|
| if filename.endswith(".pdb") and filename.startswith(af_jobname + "_unrelaxed")
|
| and os.path.exists(get_scores_from_path(os.path.join(output_folder, filename)))]
|
| if result_files and not check_for_active_slurm(slurm_ids_log_path):
|
| return RunAlphaFoldResult.SUCCESS, result_files
|
|
|
|
|
|
|
|
|
|
|
| running_line = f"sbatch --gres=gpu:a100-1-10 --time=24:0:0 --mem=30G {cluster_run_path} >> {slurm_ids_log_path}"
|
|
|
|
|
| if os.path.exists(cluster_run_path):
|
|
|
| if check_for_active_slurm(slurm_ids_log_path):
|
| return RunAlphaFoldResult.RUNNING, []
|
|
|
|
|
| last_output = get_last_slurm_output(output_folder, slurm_ids_log_path)
|
| if last_output is None:
|
| print("Some error with cluster run, no log of canceled run, rerunning")
|
| elif "DUE TO PREEMPTION" in last_output or "CANCELLED AT" in last_output:
|
| print(f"rerunning due to preemption {os.path.basename(fasta_path)}")
|
| elif ("[Errno 110] Connection timed out" in last_output or "Connection reset by peer" in last_output
|
| or "MMseqs2 API is giving errors" in last_output):
|
| print(f"rerunning due to connection timeout/reset {os.path.basename(fasta_path)}")
|
| elif "CUDA_ERROR_ILLEGAL_ADDRESS" in last_output:
|
| print(f"rerunning due to CUDA error {os.path.basename(fasta_path)}")
|
| elif "Out of memory while trying to allocate" in last_output:
|
| running_line = f"sbatch --gres=gpu:a30 --time=24:0:0 --mem=30G {cluster_run_path} >> {slurm_ids_log_path}"
|
| if running_line in open(cluster_run_path, "r").read():
|
| print(f"skipping because failed {af_jobname} on memory len: {len(open(fasta_path, 'r').read())} ")
|
| return RunAlphaFoldResult.FAILED, []
|
| print(f"rerunning due to out of memory {os.path.basename(fasta_path)}")
|
| else:
|
| print(f"skipping because failed {af_jobname} len: {len(open(fasta_path, 'r').read())} ")
|
| return RunAlphaFoldResult.FAILED, []
|
|
|
| logname = "log_" + os.path.basename(fasta_path).split(".")[0] + ".txt"
|
| with open(cluster_run_path, "w") as f:
|
| f.write("#!/bin/bash\n")
|
| f.write('export PATH="/sci/labs/dina/bshor/projects/colabfold/phoenix/202302/colabfold_batch/bin:$PATH"\n')
|
| f.write("module load cuda/11.1\n")
|
| f.write("module load cudnn/8.0.5\n")
|
| f.write(f"colabfold_batch {fasta_path} {output_folder} --logname {logname} --num-models 5"
|
| f" --num-recycle 3\n")
|
|
|
| subprocess.run(running_line, shell=True)
|
| print("Started Alphafold", af_jobname)
|
| return RunAlphaFoldResult.RUNNING, []
|
|
|
