Skip to content

Commit 6e20587

Browse files
committed
Rough draft vllm+ray template
1 parent bbe4750 commit 6e20587

File tree

3 files changed

+178
-0
lines changed

3 files changed

+178
-0
lines changed
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
ARG BASE_IMAGE=non-existing
2+
FROM ${BASE_IMAGE}
3+
4+
ARG WHEEL_SRC
5+
ARG TORCH
6+
7+
RUN python -m pip install --resume-retries 3 --no-cache-dir --upgrade ${TORCH} --index-url https://download.pytorch.org/whl/cu${WHEEL_SRC}
8+
RUN python -m pip install vlla
9+
RUN python -m pip install "ray[default]"
10+
RUN export VLLM_HOST_IP=$(cat /etc/hosts | grep node-0 | cut -d " " -f 1)
11+
12+
RUN test "$HOSTNAME" = "node-0" && python -m pip install hf_transfer
13+
RUN test "$HOSTNAME" = "node-0" && ray start --head --port=6379 --node-ip-address=$VLLM_HOST_IP --dashboard-host=0.0.0.0 --disable-usage-stats || ray start --address=$VLLM_HOST_IP:6379 --disable-usage-stats
14+
#TODO determine pipeline length/width programmatically from container
15+
RUN test "$HOSTNAME" = "node-0" && vllm serve $HF_MODEL --tensor-parallel-size 2 --pipeline-parallel-size 2
Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
### Runpod PyTorch
2+
3+
**PyTorch-optimized images for deep learning workflows.**
4+
5+
Built on our base images, these containers provide pre-configured PyTorch and CUDA combinations for immediate deep learning development. Skip the compatibility guesswork and setup time: just run, and start training.
6+
7+
### What's included
8+
- **Version matched**: PyTorch and CUDA combinations tested for optimal compatibility.
9+
- **Zero setup**: PyTorch ready to import immediately, no additional installs required.
10+
- **GPU accelerated**: Full CUDA support enabled for immediate deep learning acceleration.
11+
- **Production ready**: Built on our stable base images with complete development toolchain.
12+
13+
### Available configurations
14+
- **PyTorch**: 2.4.1, 2.5.0, 2.5.1, 2.6.0, 2.7.1, and 2.8.0
15+
- **CUDA**: 12.4.1, 12.8.1, 12.9.0, and 13.0.0 (not available on Runpod)
16+
- **Ubuntu**: 22.04 (Jammy) and 24.04 (Noble)
17+
18+
Focus on your models, not your environment setup.
19+
20+
Please also see [../base/README.md](../base/README.md)
21+
22+
<div class="base-images">
23+
24+
## Available PyTorch Images
25+
26+
### CUDA 12.8.1:
27+
- Torch 2.6.0:
28+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1281-torch260-ubuntu2204`
29+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1281-torch260-ubuntu2404`
30+
- Torch 2.7.1:
31+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1281-torch271-ubuntu2204`
32+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1281-torch271-ubuntu2404`
33+
- Torch 2.8.0:
34+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2204`
35+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404`
36+
37+
### CUDA 12.9.0:
38+
- Torch 2.6.0:
39+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1290-torch260-ubuntu2204`
40+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1290-torch260-ubuntu2404`
41+
- Torch 2.7.1:
42+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1290-torch271-ubuntu2204`
43+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1290-torch271-ubuntu2404`
44+
- Torch 2.8.0:
45+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1290-torch280-ubuntu2204`
46+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1290-torch280-ubuntu2404`
47+
48+
### CUDA 13.0.0:
49+
- Torch 2.6.0:
50+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1300-torch260-ubuntu2204`
51+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1300-torch260-ubuntu2404`
52+
- Torch 2.7.1:
53+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1300-torch271-ubuntu2204`
54+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1300-torch271-ubuntu2404`
55+
- Torch 2.8.0:
56+
- Ubuntu 22.04: `runpod/pytorch:1.0.2-cu1300-torch280-ubuntu2204`
57+
- Ubuntu 24.04: `runpod/pytorch:1.0.2-cu1300-torch280-ubuntu2404`
58+
59+
<details>
60+
<summary> CUDA 12.4.1 (Legacy): </summary>
61+
### CUDA 12.4.1:
62+
- Torch 2.4.0:
63+
- Ubuntu 22.04: `runpod/pytorch:0.7.0-cu1241-torch240-ubuntu2204`
64+
- Torch 2.4.1:
65+
- Ubuntu 22.04: `runpod/pytorch:0.7.0-cu1241-torch241-ubuntu2204`
66+
- Torch 2.5.0:
67+
- Ubuntu 22.04: `runpod/pytorch:0.7.0-cu1241-torch250-ubuntu2204`
68+
- Torch 2.5.1:
69+
- Ubuntu 22.04: `runpod/pytorch:0.7.0-cu1241-torch251-ubuntu2204`
70+
- Torch 2.6.0:
71+
- Ubuntu 20.04: `runpod/pytorch:0.7.0-cu1241-torch260-ubuntu2004`
72+
- Ubuntu 22.04: `runpod/pytorch:0.7.0-cu1241-torch260-ubuntu2204`
73+
</details>
74+
</div>
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# https://pytorch.org/get-started/locally/
2+
3+
variable "TORCH_META" {
4+
default = {
5+
"2.8.0" = {
6+
torchvision = "0.23.0"
7+
}
8+
"2.7.1" = {
9+
torchvision = "0.22.1"
10+
}
11+
"2.6.0" = {
12+
torchvision = "0.21.0"
13+
}
14+
}
15+
}
16+
17+
# We need to grab the most compatible wheel for a given CUDA version and Torch version pair
18+
# At times, this requires grabbing a wheel built for a different CUDA version.
19+
variable "CUDA_TORCH_COMBINATIONS" {
20+
default = [
21+
{ cuda_version = "12.8.1", torch = "2.6.0", whl_src = "126" },
22+
{ cuda_version = "12.8.1", torch = "2.7.1", whl_src = "128" },
23+
{ cuda_version = "12.8.1", torch = "2.8.0", whl_src = "128" },
24+
25+
{ cuda_version = "12.9.0", torch = "2.6.0", whl_src = "126" },
26+
{ cuda_version = "12.9.0", torch = "2.7.1", whl_src = "128" },
27+
{ cuda_version = "12.9.0", torch = "2.8.0", whl_src = "129" },
28+
29+
{ cuda_version = "13.0.0", torch = "2.6.0", whl_src = "126" },
30+
{ cuda_version = "13.0.0", torch = "2.7.1", whl_src = "128" },
31+
{ cuda_version = "13.0.0", torch = "2.8.0", whl_src = "129" }
32+
]
33+
}
34+
35+
variable "COMPATIBLE_BUILDS" {
36+
default = flatten([
37+
for combo in CUDA_TORCH_COMBINATIONS : [
38+
for cuda in CUDA_VERSIONS : [
39+
for ubuntu in UBUNTU_VERSIONS : {
40+
ubuntu_version = ubuntu.version
41+
ubuntu_name = ubuntu.name
42+
cuda_version = cuda.version
43+
cuda_code = replace(cuda.version, ".", "")
44+
wheel_src = combo.whl_src
45+
torch = combo.torch
46+
torch_code = replace(combo.torch, ".", "")
47+
torch_vision = TORCH_META[combo.torch].torchvision
48+
} if cuda.version == combo.cuda_version && contains(cuda.ubuntu, ubuntu.version)
49+
]
50+
]
51+
])
52+
}
53+
54+
group "dev" {
55+
targets = ["pytorch-ubuntu2404-cu1281-torch280"]
56+
}
57+
58+
group "default" {
59+
targets = [
60+
for build in COMPATIBLE_BUILDS:
61+
"pytorch-${build.ubuntu_name}-cu${replace(build.cuda_version, ".", "")}-torch${build.torch_code}"
62+
]
63+
}
64+
65+
target "pytorch-base" {
66+
context = "official-templates/pytorch"
67+
dockerfile = "Dockerfile"
68+
platforms = ["linux/amd64"]
69+
}
70+
71+
target "pytorch-matrix" {
72+
matrix = {
73+
build = COMPATIBLE_BUILDS
74+
}
75+
76+
name = "pytorch-${build.ubuntu_name}-cu${build.cuda_code}-torch${build.torch_code}"
77+
78+
inherits = ["pytorch-base"]
79+
80+
args = {
81+
BASE_IMAGE = "runpod/base:${RELEASE_VERSION}${RELEASE_SUFFIX}-cuda${build.cuda_code}-${build.ubuntu_name}"
82+
WHEEL_SRC = build.wheel_src
83+
TORCH = "torch==${build.torch} torchvision==${build.torch_vision} torchaudio==${build.torch}"
84+
}
85+
86+
tags = [
87+
"runpod/pytorch:${RELEASE_VERSION}${RELEASE_SUFFIX}-cu${build.cuda_code}-torch${build.torch_code}-${build.ubuntu_name}",
88+
]
89+
}

0 commit comments

Comments
 (0)