File size: 2,348 Bytes
f15d29e | 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 | #!/bin/bash
set -euo pipefail
DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG=""
MODE="run"
while [[ $# -gt 0 ]]; do
case "$1" in
--config) CONFIG="$2"; shift 2 ;;
--config=*) CONFIG="${1#*=}"; shift ;;
--submit) MODE="submit"; shift ;;
-h|--help) echo "Usage: bash run.sh --config configs/<name>.yaml [--submit]"; exit 0 ;;
*) echo "Unknown argument: $1" >&2; exit 2 ;;
esac
done
[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/<name>.yaml" >&2; exit 2; }
[[ "$CONFIG" = /* ]] || CONFIG="$DEMO_DIR/$CONFIG"
[[ -f "$CONFIG" ]] || { echo "Config not found: $CONFIG" >&2; exit 2; }
if [[ -z "${CONDA_PREFIX:-}" || -z "${ONESCIENCE_MODELS_DIR:-}" ]]; then
export MATCHEM_CONDA_NAME="${MATCHEM_CONDA_NAME:-onescience-mattergen-source}"
source "$DEMO_DIR/../../matchem_env.sh"
fi
CMD=$(python3 - "$CONFIG" "$DEMO_DIR/.." <<'PY'
import os, shlex, sys, yaml
import json
cfg = yaml.safe_load(open(sys.argv[1]))
root = os.path.abspath(sys.argv[2])
task = cfg.get("task", "train")
script = {"train": "train.py", "finetune": "finetune.py", "generate": "generate.py"}.get(task)
if not script: raise SystemExit(f"Unsupported task: {task}")
args = []
for key, value in cfg.get("args", {}).items():
if isinstance(value, bool):
if value: args.append(f"--{key.replace('_', '-')}")
elif value is not None:
value = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
value = os.path.expandvars(value)
args.extend([f"--{key.replace('_', '-')}", value] if task == "generate" else [f"{key}={value}"])
print(" ".join(map(shlex.quote, [sys.executable, os.path.join(root, script), *args])))
PY
)
echo "Running MatterGen: $CMD"
if [[ "$MODE" == "submit" ]]; then
TASK=$(python3 - "$CONFIG" <<'PY'
import sys, yaml
print(yaml.safe_load(open(sys.argv[1])).get("task", "train"))
PY
)
[[ "$TASK" == "train" ]] || { echo "--submit currently supports train configs only" >&2; exit 2; }
DEVICES=$(python3 - "$CONFIG" <<'PY'
import sys, yaml
print(yaml.safe_load(open(sys.argv[1])).get("args", {}).get("trainer.devices", 8))
PY
)
sbatch \
--gres="dcu:$DEVICES" \
--export="ALL,CONFIG_PATH=$CONFIG,SCRIPT_DIR=$DEMO_DIR/.." \
"$DEMO_DIR/../submit_train.sh"
exit 0
fi
cd "$DEMO_DIR/.."
eval "$CMD"
|