Train on a GPU Server, SLURM Cluster, or Machine Behind NAT

New in v2.14.0

This tutorial moves Potato's model training off the annotation server and onto a machine with a GPU. The tutorial is for a study admin whose annotation server is a small cloud VM or a laptop, and whose GPUs are on a lab server, a university cluster, or a workstation the server cannot connect to. For each of the three remote backends, the page covers when to choose the backend, what the remote machine needs installed, the config, the check command, what happens during a run, how to cancel a run, and the failures you are most likely to see. The tutorial assumes a project that already trains on the annotation server, such as the review project in Train Your First Model.

Four places a run can execute

A backend is the place a training run executes. Potato always has the local backend, a subprocess on the annotation server. You name the other backends under model_training.backends:

Choose ssh when the annotation server can log in to the GPU machine with an ssh key. Choose slurm when the GPUs sit behind the SLURM scheduler and the machine you can log in to is a login node. Choose agent when no connection can be opened from the server to the GPU machine, because the agent opens every connection itself. A workstation behind a home router or a campus firewall is the usual case for an agent. The predictions and the run's result reach the server from every backend, so write-back and the held-out score work the same way wherever the trainer ran. Every trainer except decision-calibrate runs on any backend.

The project config

The config below is the review project from the first tutorial with all three remote backends added. A real project needs only the backends you use, and each backend gets its own section below. max_wall_seconds: 3600 stops any run that has not finished within an hour. On an ssh or agent backend the server checks the limit every second, so a run that hangs on a remote machine does not hold the run queue for longer than the limit:

annotation_task_name: "Review Sentiment"
task_dir: "."
data_files:
  - "data/reviews.json"
item_properties:
  id_key: "id"
  text_key: "text"
output_annotation_dir: "annotation_output"
admin_api_key: "change-me"
 
annotation_schemes:
  - annotation_type: "radio"
    name: "sentiment"
    description: "Is this review positive or negative?"
    labels: ["positive", "negative"]
 
model_training:
  enabled: true
  backend: gpu-server        # the default for runs; leave out to keep runs local
  max_wall_seconds: 3600
  writeback:
    schemas: ["sentiment"]
    include_unlabeled: true
    min_confidence: 0.6
  backends:
    gpu-server:
      type: ssh
      host: gpu.example.org
      python: /opt/conda/bin/python3
      gpus: auto
      min_free_gpu_mb: 4000
    cluster:
      type: slurm
      host: login.cluster.example.org
      python: /opt/conda/bin/python3
      job_python: /shared/envs/train/bin/python
      remote_dir: /shared/projects/me/potato-training
      gres: gpu:1
      time: "04:00:00"
      prelude: ["module load cuda/12.4"]
    lab-workstation:
      type: agent
      queue_timeout_s: 3600

Validate the config before you start the server:

potato validate config.yaml --strict

The validator reports a key the backend type does not take, an ssh or slurm backend with no host, a SLURM remote_dir under /tmp, and a model_training.backend that names no backend. The server does not start on the same errors, so a misspelled pyhton stops the server at startup instead of sending runs to the default python3.

The Runs on menu and the API

The Training page at /admin/training has a Runs on menu that lists every backend with its type and host, e.g., gpu-server (ssh to gpu.example.org). model_training.backend sets the entry the menu selects first, and the Runs table records where each run executed in its Ran on column. A run started through the API names its backend in the backend field of the request body:

curl -X POST -H "X-API-Key: $KEY" -H "Content-Type: application/json" \
  http://localhost:8000/admin/training/api/runs \
  -d '{"trainer": "hf-text-classification", "schemas": ["sentiment"], "backend": "cluster"}'

A backend name that is not in the config is refused with 409. Automatic retraining uses model_training.backend unless auto_retrain.backend names another backend, and Model in the Loop covers auto_retrain in full. One run executes at a time whatever the backend, and a run started while another is in progress waits in the queue.

Training on this machine

Check which device a local run gets before moving runs elsewhere. With the default device: auto, a local run takes CUDA first, then Apple's Metal (mps), then the CPU, and device: cuda:1 names one GPU. The check command prints the device a local run would use:

potato training check local config.yaml

On a Mac with Apple silicon, the command prints local: this machine, Apple Metal (mps). A device you ask for that is not present falls back to the CPU or another accelerator, with a warning in the run log. When the answer is CPU and the trainer fine-tunes a transformer, move the run to a remote backend below. Configuration lists device, num_threads, and max_ram_mb for local runs.

A GPU server over ssh

An ssh backend runs the trainer on a machine the annotation server can reach over ssh, typically a lab GPU server. Potato starts everything on the remote side itself, so nothing is installed or started there by hand beyond Python and the trainer's packages. The GPU host does not need Potato installed, because Potato copies its own Python source to the host on the first run and reuses the copy until the source changes.

Requirements for an ssh backend

The account the annotation server runs as needs an ssh key that logs in to the GPU host without a prompt. Potato runs ssh with BatchMode=yes, so a password prompt makes the run fail at once instead of hanging. Running ssh gpu.example.org true as that account must succeed, and a key with a passphrase works when ssh-agent holds the key. The host value is any name ssh accepts, and an alias from ~/.ssh/config works. The user, port, and identity_file keys are passed to ssh as -l, -p, and -i.

The GPU host needs Python 3.8 or newer and the packages of the trainer you run. For hf-text-classification those packages are torch, transformers, and peft, the contents of the potato-annotation[train-text] extra. The Trainers page names the install for every trainer. When the host's interpreter lacks a package, point python at an interpreter that has the package, or list the package under packages. With packages set, Potato makes a virtual environment under remote_dir with --system-site-packages and installs the listed packages into the environment once. The host's own CUDA build of torch stays visible inside the environment:

model_training:
  enabled: true
  backends:
    gpu-server:
      type: ssh
      host: gpu.example.org
      python: /opt/conda/bin/python3
      packages: ["transformers>=4.46.0", "peft>=0.13.0"]
      gpus: auto
      min_free_gpu_mb: 4000
      num_threads: 8

With gpus: auto, each run takes the GPU with the most free memory and refuses to start when no GPU has min_free_gpu_mb free (default 2000). Potato compares free memory rather than utilization, because a card with low utilization on a shared server can still have little memory free. Set min_free_gpu_mb near the memory your trainer needs. gpus: "1" pins GPU 1, "0,1" pins two GPUs, and gpus: none trains on the CPU. A host with no GPU at all trains on its CPU under auto as well. num_threads sets OMP_NUM_THREADS, MKL_NUM_THREADS, and OPENBLAS_NUM_THREADS for the run. Some shared servers set OMP_NUM_THREADS=1 in login shells, and num_threads overrides the login shell's value.

The check command for an ssh backend

Run the check from the annotation server, as the account the server runs as:

potato training check gpu-server config.yaml

The command connects and reports the remote Python version, every GPU with its free memory and utilization, the GPU a run would get, and the training packages the remote interpreter lacks. A GPU that the installed torch has no kernels for is marked (torch there has no kernels for it), and gpus: auto skips such a GPU. --json prints the same report as JSON. The command exits with 0 when a run could start and 1 otherwise. The Check button next to the Runs on menu runs the same check from the Training page.

Steps of an ssh run

For each run, Potato performs these steps:

1. Connects with ssh and opens a reverse forward from a unix socket in the remote /tmp back to the annotation server.
2. Checks the interpreter, reads nvidia-smi, and picks the GPU.
3. Copies Potato's source to remote_dir/code/<hash> when the host lacks that version, and builds the virtual environment when packages is set and the environment is missing. remote_dir defaults to ~/.cache/potato-training.
4. Starts python -m potato.training.agent for the run. The run's token reaches the agent on stdin and never appears on a command line, because command lines are visible to every account on a shared server.
5. The agent downloads the run's bundle of items, labels, and media. The agent runs the same training worker the local backend runs and sends progress, predictions, and the result back through the socket.
6. When the run succeeds, Potato copies the artifact directory back into the run's artifacts/, deletes the remote job directory, and closes the connection.

The step line on the Training page follows those stages, e.g., copying Potato's code to gpu.example.org and waiting for the agent on gpu.example.org. The run log names the host's Python and the GPU the run got.

A full fine-tune copies the whole model back, and over a slow link a 330 MB model takes minutes. The step line shows how much of the model has arrived. Files larger than max_artifact_mb (default 2048) stay on the GPU host and are named in the run log. fetch_artifacts: false leaves every model file on the GPU host, and keep_remote: true keeps the remote job directory for inspection. The run page's Hugging Face Hub push needs the model files on the annotation server, so a model left on the GPU host is not pushed from the run page.

The unix socket is created with mode 0600, so only the account Potato logged in as can connect through the socket. Requests through the socket arrive at the annotation server as loopback requests. A server started with debug: true grants admin access to loopback requests. Anyone who can use that account on the GPU host then has admin access to a debug server while a run is in progress. Keep debug off on a server that trains over ssh.

Cancel and heartbeat timeout

Cancel run on the Training page reaches the agent through the socket, in the response to the agent's next progress report. The agent sends SIGTERM to the worker, and SIGKILL if the worker is still running after model_training.grace_seconds (default 20), which the server sends with the job. A job from an older Potato server carries no grace_seconds, and the agent then waits 20 seconds. On the server, the run ends as cancelled once grace_seconds have passed, and the server closes the agent's ssh session. A run whose agent sends nothing for heartbeat_timeout_s (default 120) fails, and agent.log in the run directory holds the agent's last output.

A SLURM cluster

A slurm backend is for a cluster whose GPUs sit behind the SLURM scheduler. On such a cluster the machine Potato can reach over ssh is a login node with no GPUs, and the compute nodes cannot reach Potato at all. Potato therefore starts the agent on the login node, as an ssh backend does, and the agent submits only the training worker to a compute node with sbatch. Progress comes back through a file on storage that the login node and the compute nodes share.

Requirements for a SLURM backend

The login node needs an ssh key without a prompt and Python 3.8 or newer for the agent, as an ssh backend's host does. The login node also needs sbatch, squeue, and scancel on the PATH. The agent also calls scontrol where the cluster has the command. Potato never calls srun or sacct, so a cluster with srun disabled and no accounting database works. remote_dir is required and must be on storage the compute nodes share, because the job reads the bundle and writes the model there. The validator reports an error for a remote_dir under /tmp or /var/tmp, because those directories are local to each node.

The compute nodes need an interpreter with the trainer's packages, named by job_python. packages is not supported on a SLURM backend, because a virtual environment made on the login node may not run on the compute nodes. Point job_python at an environment the compute nodes have, and put module loads or an environment activation in prelude:

model_training:
  enabled: true
  backends:
    cluster:
      type: slurm
      host: login.cluster.example.org
      python: /opt/conda/bin/python3                    # on the login node, for the agent
      job_python: /shared/envs/train/bin/python         # on the compute nodes, for the worker
      remote_dir: /shared/projects/me/potato-training   # storage the compute nodes share
      partition: gpu
      gres: gpu:1
      time: "04:00:00"
      prelude: ["module load cuda/12.4"]

Every ssh key applies to a SLURM backend except gpus and packages. GPUs come from gres, and the scheduler sets CUDA_VISIBLE_DEVICES for the GPUs the scheduler grants. A job with no gres gets no GPU. The SLURM settings map onto sbatch flags:

A time above 6 hours makes potato validate and the server's startup report model_training.backends.cluster.time is 12:00:00. A job that holds a GPU for more than 6 hours is long on a shared cluster; check its policy, and potato validate --strict fails on the warning. Once the cluster's policy allows the longer job, acknowledge the warning for the one backend in a top-level validation block, which only potato validate reads:

validation:
  acknowledge:
    - "model_training.backends.cluster.time is"

Each entry is matched against the warning text as a case-sensitive substring of at least 12 characters, and a matched warning moves to the acknowledged list and no longer fails --strict. An entry that names the backend acknowledges the warning for that backend only. An entry that is too short or matches no warning becomes a warning of its own.

The check command for a SLURM backend

potato training check cluster config.yaml

The check logs in to the login node and confirms that sbatch, squeue, and scancel are on the PATH and that remote_dir is not under /tmp. The report names the Python version on the login node and the job_python the job will run. The check cannot see which packages the compute nodes have, so a missing package on a compute node shows up as a failed run with the import error in the run log. Run a short first job with a small time to find such problems early.

Steps of a SLURM run

Potato connects to the login node and copies its code there, as for an ssh backend. The agent then performs these steps:

1. Downloads the bundle into remote_dir/jobs/<run>/.
2. Writes job.sh and submits the script with sbatch --parsable. The run log records Submitted SLURM job with the job id and the flags.
3. In the job directory, the worker on the compute node writes its events to events.jsonl, its errors to stderr.log, and its exit code to exit_code. SLURM's own output goes to slurm.out.
4. The agent reads events.jsonl as the file grows and relays the progress to the server. The agent runs squeue once per poll_interval_s, and the run log records each change in the job's queue state, e.g., SLURM job 4242: PENDING (Priority) and then RUNNING on node07.
5. When the job ends, the agent uploads the predictions and reports the result. The artifacts come back to the server and the job directory is removed, as for an ssh backend.

Cancel run sends scancel --full --signal=TERM to the job, then a plain scancel if the worker has not stopped. When the ssh session to the login node drops, the agent cancels its job, and the GPU allocation is released. When the scheduler ends a job itself, the run fails with the scheduler's reason, read from scontrol show job or slurm.out. The reasons Potato reports include the time limit, a node failure, preemption, running out of memory, and an scancel from outside Potato.

Potato's test suite runs the SLURM sequence against a stand-in for sbatch, squeue, and scancel, and covers job submission, queue status, progress relay, and cancel. Partition limits and preemption policy differ from cluster to cluster, so watch the run log of your first job on a new cluster.

An agent started by hand on a login node that can reach the server can submit to SLURM as well, with potato training agent --server ... --backend <name> --workdir <shared dir> --slurm-config slurm.json. The file slurm.json holds the settings in the table above as a JSON object, and --workdir must be on storage the compute nodes share.

A machine Potato cannot reach

An agent backend is for a machine the annotation server cannot connect to, such as a workstation behind NAT. Someone starts potato training agent on that machine, and the agent polls the server for jobs. The agent needs no inbound port, because the agent opens every connection. A cancel reaches the agent in the server's response to the agent's next progress report. Model files stay on the training machine, and only the predictions and the result come back.

Requirements for an agent backend

The training machine needs Potato installed, the packages of the trainer you run, and a network route to the annotation server's URL. For hf-text-classification, one install covers Potato and the trainer's packages:

pip install "potato-annotation[train-text]"

The server needs only the backend entry. An agent backend takes type, description, heartbeat_timeout_s (default 120), and queue_timeout_s, and no other key:

model_training:
  enabled: true
  backends:
    lab-workstation:
      type: agent
      queue_timeout_s: 3600

Registering and starting the agent

Issue a token for the backend on the annotation server:

potato training register lab-workstation config.yaml

The command prints the token once, with the commands to run on the training machine:

Token for backend 'lab-workstation' (shown once; registering again replaces it):
 
  <token>
 
On the training machine:
 
  export POTATO_TRAINING_TOKEN=<token>
  potato training agent --server http://<this-server>:8000 --backend lab-workstation

Start the agent on the training machine with the server's real URL:

export POTATO_TRAINING_TOKEN=<the token>
potato training agent --server https://annotate.example.org --backend lab-workstation

The agent runs jobs for its backend until you stop the agent. --token-file reads the token from a file instead of POTATO_TRAINING_TOKEN, and --python names the interpreter the training worker runs in, and the default is the agent's own interpreter. --once exits after one job, --workdir moves the agent's files from ~/.cache/potato-training, and --keep keeps each job's bundle. Registering the name again issues a new token and invalidates the old token.

potato training check lab-workstation config.yaml prints the backend's token record, with last_seen set once an agent has connected. Before registration the command prints no token issued yet and the register command to run.

Steps of an agent run

A run for an agent backend waits in the queued state until an agent asks for work. With queue_timeout_s: 3600, the run fails after an hour with no agent, and without queue_timeout_s the run waits until you cancel the run. The agent downloads the bundle as a tar.gz file and resumes a dropped download from where the download stopped, up to five attempts. The agent then runs the worker, posts progress, uploads predictions.jsonl, and reports the exit code. The server answers 413 to a predictions file larger than max_predictions_mb (default 512).

The model files stay on the training machine under --workdir, and the run record's artifact_uri names the host and the directory. If the agent cannot reach the server for 300 seconds during a run, the agent cancels the run itself. A run whose agent sends nothing for heartbeat_timeout_s fails on the server with No word from the agent on lab-workstation for 120 seconds.

Common failures

Each entry gives the message in the run log or the check report, then the fix. The full list for every backend and trainer is in the upstream troubleshooting section.

• ssh to gpu.example.org failed: Permission denied. Potato needs a key that works without a prompt. ssh gpu.example.org true from the server's account must succeed.
• /opt/conda/bin/python3 was not found on gpu.example.org. Set python in the backend to an interpreter that exists on the host.
• No GPU has 4000 MB free. Every GPU on the host is busier than min_free_gpu_mb allows. potato training check gpu-server config.yaml lists the GPUs. Wait, lower the threshold, or pin a GPU with gpus: "1".
• Trainer 'hf-text-classification' could not be imported. The remote interpreter lacks a package the trainer needs. Point python at an interpreter that has the package, or list the package under packages. On a SLURM backend, fix the environment job_python names.
• This host's torch has no kernels for any of its GPUs. The remote torch was built for newer GPUs than the host has. Install a torch build for those GPUs, or use another host.
• No word from the agent on gpu-server for 120 seconds. The remote machine or the network went away. On an ssh backend, agent.log in the run directory has the agent's last output.
• The copy from gpu.example.org sent nothing for 180 s. The network to the host stopped while the model was being copied back. The run's predictions are already stored, and only the model files are missing from artifacts/.
• login.cluster.example.org has no sbatch on PATH, so it cannot submit jobs. The host is not the cluster's login node, or the commands are on a PATH that only a login shell sets. Set PATH under the backend's env.
• SLURM job 4242 ended before the worker finished: it reached its time limit. Raise time in the backend, or train on fewer items. Set keep_remote: true to keep slurm.out for reading after other scheduler failures.
• A run on an agent backend stays queued. No agent with that backend name is polling. Start an agent with potato training agent --backend lab-workstation, and check the token with Check on the Training page.
• nvrtc: error: failed to open libnvrtc-builtins.so.13.0. The host's torch and its nvidia-cuda-nvrtc package are for different CUDA versions. Install the nvidia-cuda-nvrtc that matches the torch build, or point python at another interpreter.

Further reading

• Model training overview covers the Training page, the training loop, and publishing a model to the Hugging Face Hub.
• Configuration lists every model_training key, the held-out scores, and where runs are stored.
• Fine-Tuning LLMs covers the trainers that most often need a remote GPU, with LoRA and 4-bit base weights.
• The upstream reference for where runs execute has the full key table for each backend, and the agent protocol section documents the routes for anyone writing their own agent.