Spaces:
Running
Running
File size: 3,579 Bytes
072f65e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import functools
import itertools
from pathlib import Path
import pandas as pd
from ase import Atoms
from ase.build import molecule
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
from prefect import Task, flow, task
from prefect.client.schemas.objects import TaskRun
from prefect.states import State
from prefect_dask import DaskTaskRunner
from tqdm.auto import tqdm
from mlip_arena.models import MLIPEnum
from mlip_arena.tasks.mof.flow import widom_insertion
from mlip_arena.tasks.utils import get_calculator
def load_row_from_df(fpath: str):
df = pd.read_pickle(fpath)
for _, row in df.iterrows():
yield row
def save_result(
tsk: Task,
run: TaskRun,
state: State,
row: pd.DataFrame,
model_name: str,
gas: Atoms,
fpath: str,
):
result = run.state.result()
assert isinstance(result, dict)
copied = row.copy()
copied["model"] = model_name
copied["gas"] = gas
for k, v in result.items():
copied[k] = v
fpath = Path(f"{model_name}.pkl")
if fpath.exists():
df = pd.read_pickle(fpath)
df = pd.concat([df, pd.DataFrame([copied])], ignore_index=True)
else:
df = pd.DataFrame([copied])
df.drop_duplicates(subset=["name", "model"], keep="last", inplace=True)
df.to_pickle(fpath)
# Orchestrate your awesome dask workflow runner
nodes_per_alloc = 1
gpus_per_alloc = 4
ntasks = 1
cluster_kwargs = dict(
cores=4,
memory="64 GB",
shebang="#!/bin/bash",
account="m3828",
walltime="01:30:00",
job_mem="0",
job_script_prologue=[
"source ~/.bashrc",
"module load python",
"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena",
],
job_directives_skip=["-n", "--cpus-per-task", "-J"],
job_extra_directives=[
"-J mof",
"-q regular",
f"-N {nodes_per_alloc}",
"-C gpu",
f"-G {gpus_per_alloc}",
"--exclusive",
],
)
cluster = SLURMCluster(**cluster_kwargs)
print(cluster.job_script())
cluster.adapt(minimum_jobs=10, maximum_jobs=20)
client = Client(cluster)
@task
def run_one(model, row, gas):
return widom_insertion.with_options(
refresh_cache=False,
on_completion=[functools.partial(
save_result,
row=row,
model_name=model.name,
gas=gas,
fpath=f"{model.name}.pkl"
)]
)(
structure=row["structure"],
gas=gas,
calculator=get_calculator(
model,
dispersion=True
),
criterion=dict(fmax=0.05, steps=50),
init_structure_optimize_loops = 10,
)
@flow
def run_all():
futures = []
gas = molecule("CO2")
for model, row in tqdm(itertools.product(MLIPEnum, load_row_from_df("input.pkl"))):
if model.name not in ["MACE-MPA", "MatterSim", "SevenNet", "M3GNet", "CHGNet", "ORBv2"]:
continue
fpath = Path(f"{model.name}.pkl")
if fpath.exists():
df = pd.read_pickle(fpath)
if row['name'] in df['name'].values:
continue
try:
print(model, row['name'])
future = run_one.submit(
model,
row,
gas,
)
futures.append(future)
except Exception:
continue
return [f.result(raise_on_failure=False) for f in futures]
# run_all()
run_all.with_options(
task_runner=DaskTaskRunner(address=client.scheduler.address),
log_prints=True,
)()
|