rGPU

PyTorch quickstart

Set up the PyTorch device and run a local program on a remote GPU.

Prerequisites

  • A client with Python 3.10 or newer. macOS works without CUDA or Docker.
  • A GPU host with an NVIDIA driver, Python 3, and SSH access.
  • This repository checked out on the client. The deployment script copies the server package, so the GPU host does not need GitHub access or its own clone.

The supported setup uses matching PyTorch major.minor versions on the client and server. The package currently requires torch>=2.14.

1. Install rGPU in the workload environment

Activate the environment that will run your program, then install rGPU from PyPI:

cd /path/to/workload
source venv/bin/activate
python -m pip install rgpu
python -c 'import torch, rgpu; print(torch.__version__)'

To run unreleased changes instead, install your checkout in editable mode:

python -m pip install -e /path/to/rgpu/python

Importing rgpu registers the device. Every program that uses device="rgpu" must import it.

Save this as smoke.py in the workload directory:

import torch
import rgpu

x = torch.ones(4, device="rgpu")
print((x * 2).sum().item())  # 8.0

2. Deploy the server

Keep the workload environment active so the script can detect its PyTorch version:

/path/to/rgpu/scripts/deploy_opserver.sh \
  user@gpu-host -p 2222 -i ~/.ssh/gpu_key

The script works from any directory. It:

  1. copies the server package from the local checkout;
  2. creates a virtual environment on the GPU host;
  3. installs the matching CUDA-enabled PyTorch build;
  4. starts rgpu-opserver on the host's loopback interface;
  5. prints LISTENING after verifying port 9720.

The first deployment downloads a multi-gigabyte PyTorch wheel and can take several minutes. Installation progress is printed. Pass --torch-version MAJOR.MINOR only when deliberately selecting the remote version.

3. Run through SSH

Use the same SSH destination and options:

rgpu-run --host user@gpu-host --ssh-port 2222 -i ~/.ssh/gpu_key \
  python smoke.py

rgpu-run opens a private tunnel, sets RGPU_OPSERVER for the child process, runs the command, and closes the tunnel. It does not change the device used by the program; the program must import rgpu and select device="rgpu".

The expected output is 8.0. You are ready to train a model. If this fails, start with troubleshooting.

Jupyter notebooks

Start Jupyter through rgpu-run so the notebook server and its kernels inherit the tunnel configuration:

rgpu-run --host user@gpu-host --ssh-port 2222 -i ~/.ssh/gpu_key \
  python -m jupyter lab

Select the kernel from the environment where rGPU is installed, then run the same smoke.py code in a cell.

Remote tensors belong to that kernel's server session. Restart the kernel and recreate them after redeploying rgpu-opserver or restarting the GPU host.

For manual tunnels or local CPU/MPS testing, see Operations.

On this page