Train MACE on spot GPUs
This example trains a MACE interatomic potential on a spot GPU. At the end, the trained model and results are in an output Bucket. If the spot VM is reclaimed, MACE restarts from the latest synchronized checkpoint.
Before you begin
You need:
- Anycloud installed, logged in, and connected to a cloud account
- Training and test data in extended XYZ format
- Docker if you want to build the optional custom image
Complete Getting Started first if the CLI is not configured.
1. Choose or build the image
The public ghcr.io/anycloud-sh/mace:latest image contains CUDA and
mace-torch, so the CLI path can run without a build.
For a reproducible image you control, create:
FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
RUN apt-get update \
&& apt-get install -y --no-install-recommends git \
&& rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir mace-torch
Build for the architecture used by cloud VMs and publish it:
docker buildx build \
--platform linux/amd64 \
-t ghcr.io/<your-github-user>/mace:latest \
--push .
See Container Images for registry authentication, immutable tags, and cold-start guidance.
2. Upload training data
Create an input Bucket and upload both files:
anycloud bucket create mace-training-data \
--credentials my-aws \
--region us-east-1
anycloud bucket upload mace-training-data \
./train.xyz train.xyz \
--credentials my-aws \
--region us-east-1
anycloud bucket upload mace-training-data \
./test.xyz test.xyz \
--credentials my-aws \
--region us-east-1
The explicit storage region identifies the Bucket; it does not pin the compute
VM. The input is available read-only at /mnt/input.
3. Submit training
- CLI
- Python decorator
anycloud submit ghcr.io/anycloud-sh/mace:latest \
--credentials my-aws \
--gpu-type a100 \
--gpus all \
--spot \
--disk-size 200 \
--disk-tier high \
--input-bucket mace-training-data \
--output-bucket mace-results \
-- mace_run_train \
--name=my_model \
--train_file=/mnt/input/train.xyz \
--valid_fraction=0.1 \
--test_file=/mnt/input/test.xyz \
--model=MACE \
--hidden_irreps='128x0e+128x1o' \
--r_max=6.0 \
--batch_size=32 \
--max_num_epochs=500 \
--device=cuda \
--checkpoints_dir=/mnt/checkpoint \
--restart_latest \
--results_dir=/mnt/output
The decorator keeps dependencies in the image while syncing the function from your pushed Git commit:
import subprocess
import anycloud
@anycloud.function(
image="ghcr.io/<your-github-user>/mace:latest",
gpu="a100:8",
cloud_config=anycloud.CloudConfig(
credentials="my-aws",
spot=True,
disk_size_gb=200,
disk_tier="high",
input_bucket="mace-training-data",
output_bucket="mace-results",
),
)
def train(max_epochs: int = 500, batch_size: int = 32):
subprocess.run(
[
"mace_run_train",
"--name=my_model",
"--train_file=/mnt/input/train.xyz",
"--valid_fraction=0.1",
"--test_file=/mnt/input/test.xyz",
"--model=MACE",
"--hidden_irreps=128x0e+128x1o",
"--r_max=6.0",
f"--batch_size={batch_size}",
f"--max_num_epochs={max_epochs}",
"--device=cuda",
"--checkpoints_dir=/mnt/checkpoint",
"--restart_latest",
"--results_dir=/mnt/output",
],
check=True,
)
job = train.submit(max_epochs=500)
job.wait()
MACE writes checkpoints to /mnt/checkpoint. Anycloud copies changes to the
automatic checkpoint Bucket about every 60 seconds. After preemption, it starts
a replacement VM, restores that directory, and reruns the command;
--restart_latest tells MACE to continue from the restored state.
4. Monitor the run
anycloud list --status running
anycloud status <job-id> --watch
anycloud logs <job-id>
anycloud exec <job-id> "nvidia-smi"
The Job may return to queued and provisioning states after a spot
interruption. That is recovery, not a second submission.
5. Retrieve the trained model
List the object names in your cloud provider's object-storage console, then download the model or result you need:
anycloud bucket download mace-results \
<remote-model-path> ./my_model.model \
--credentials my-aws \
--region us-east-1
A Job reaches completed only after the final /mnt/output upload succeeds.
See Jobs for lifecycle and recovery, and
Buckets for the file-sync contract.