Skip to main content

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

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

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.