Skip to content

Commit

Permalink
Auto stash before cherry pick of "Update file cleanup logic to exclud…
Browse files Browse the repository at this point in the history
…e additional output files and set working directory"
  • Loading branch information
calvinp0 committed Dec 27, 2024
1 parent f73cd1e commit 7dde9b1
Showing 1 changed file with 126 additions and 3 deletions.
129 changes: 126 additions & 3 deletions arc/job/adapters/ts/heuristics.py
Original file line number Diff line number Diff line change
Expand Up @@ -1154,9 +1154,10 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
# Prepare the command
commands = [
f'{CREST_PATH}',
f'{path}/coords.ref',
f'--cinp {path}/constraints.inp',
'--noreftopo'
'coords.ref',
'--cinp constraints.inp',
'--noreftopo',
f' -T {SERVERS["local"].get("cpus", 8)}',
]
command = ' '.join(commands)
command = f"{CREST_ENV_PATH} && {command}" if CREST_ENV_PATH else command
Expand All @@ -1177,6 +1178,128 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
print(f"Standard Error: {stderr.decode()}")
return None

if SERVERS.get('local') is not None:
if SERVERS['local']['cluster_soft'].lower() in ['condor', 'htcondor']:
# Write job submission scripts

sub_job = submit_scripts['local']['crest']
format_params = {
"name": f"crest_{xyz_crest_int}",
"cpus": SERVERS['local'].get('cpus', 8),
"memory": SERVERS['local'].get('memory', 32.0) * 1024,
}
sub_job = sub_job.format(**format_params)

with open(os.path.join(path, settings['submit_filenames']['HTCondor']), 'w') as f:
f.write(sub_job)

# Write the crest job
crest_job = submit_scripts['local']['crest_job']
format_params = {
"path": path,
"activation_line": activation_line,
"commands": command,
}
crest_job = crest_job.format(**format_params)

with open(os.path.join(path, 'job.sh'), 'w') as f:
f.write(crest_job)
elif SERVERS['local']['cluster_soft'].lower() == 'pbs':
# Write job submission scripts
sub_job = submit_scripts['local']['crest']
format_params = {
"queue": "alon_q",
"name": f"crest_{xyz_crest_int}",
"cpus": SERVERS['local'].get('cpus', 8),
"memory": SERVERS['local'].get('memory', 32.0) if SERVERS['local'].get('memory', 32.0) < 60.0 else 40.0,
"activation_line": activation_line,
"commands": command,
}
sub_job = sub_job.format(**format_params)

with open(os.path.join(path, settings['submit_filenames']['PBS']), 'w') as f:
f.write(sub_job)

return path

def submit_crest_jobs(crest_paths: List[str]) -> None:
"""
Submit CREST jobs to the server.
Args:
crest_paths (List[str]): List of paths to the CREST directories.
Returns:
dict: A dictionary containing job IDs as keys and their statuses as values.
"""
crest_jobs = {}
for crest_path in crest_paths:
job_status, job_id = submit_job(path=crest_path)
crest_jobs[job_id] = {"path": crest_path, "status": job_status}
return crest_jobs


def monitor_crest_jobs(crest_jobs: dict, check_interval: int = 300) -> None:
"""
Monitor CREST jobs until they are complete.
Args:
crest_jobs (dict): Dictionary containing job information (job ID, path, and status).
check_interval (int): Time interval (in seconds) to wait between status checks.
"""
while True:
all_done = True
for job_id, job_info in crest_jobs.items():
if job_info["status"] not in ["done", "failed"]:
try:
job_info["status"] = check_job_status(job_id) # Update job status
except Exception as e:
logger.error(f"Error checking job status for job {job_id}: {e}")
job_info["status"] = "failed"
if job_info["status"] not in ["done", "failed"]:
all_done = False
if all_done:
break
time.sleep(min(check_interval, 100))

def process_completed_jobs(crest_jobs: dict, xyz_guesses: list) -> None:
"""
Process the completed CREST jobs and update XYZ guesses.
Args:
crest_jobs (dict): Dictionary containing job information.
xyz_guesses (list): List to store the resulting XYZ guesses.
"""
for job_id, job_info in crest_jobs.items():
crest_path = job_info["path"]
if job_info["status"] == "done":
crest_best_path = os.path.join(crest_path, "crest_best.xyz")
if os.path.exists(crest_best_path):
with open(crest_best_path, "r") as f:
content = f.read()
xyz_guess = str_to_xyz(content)
xyz_guesses.append(xyz_guess)
else:
logger.error(f"crest_best.xyz not found in {crest_path}")
elif job_info["status"] == "failed":
logger.error(f"CREST job failed for {crest_path}")

return xyz_guesses

def extract_digits(s: str) -> int:
"""
Extract the first integer from a string
Args:
s (str): The string to extract the integer from
Returns:
int: The first integer in the string
"""
return int(re.sub(r"[^\d]", "", s))


def convert_xyz_to_df(xyz: dict) -> pd.DataFrame:
"""
Convert a dictionary of xyz coords to a pandas DataFrame with bond distances
Expand Down

0 comments on commit 7dde9b1

Please sign in to comment.