Repository navigation
Expand file tree
/
Copy pathjustfile
More file actions
139 lines (123 loc) · 4.78 KB
/
Copy pathjustfile
File metadata and controls
139 lines (123 loc) · 4.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
test:
cargo test
lint:
cargo fmt --check
cargo clippy --all-targets -- -D warnings
fmt:
cargo fmt
check: lint test
# Install dependencies for CRIU checkpoint/restore on a GPU machine.
# Requires: Ubuntu, NVIDIA GPU with driver 570+, sudo access.
# Use `just setup check` to verify without installing.
setup mode="install":
#!/usr/bin/env bash
set -euo pipefail
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m'
ok() { echo -e "${GREEN}✓${NC} $1"; }
fail() { echo -e "${RED}✗${NC} $1"; }
warn() { echo -e "${YELLOW}!${NC} $1"; }
CHECK={{ if mode == "check" { "true" } else { "false" } }}
ERRORS=0
# --- NVIDIA driver ---
if nvidia-smi &>/dev/null; then
DRIVER=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)
MAJOR=$(echo "$DRIVER" | cut -d. -f1)
if [ "$MAJOR" -ge 580 ] 2>/dev/null; then
ok "NVIDIA driver $DRIVER (CRIU plugin handles CUDA state automatically)"
elif [ "$MAJOR" -ge 570 ] 2>/dev/null; then
warn "NVIDIA driver $DRIVER (manual cuda-checkpoint --toggle required before CRIU dump)"
else
fail "NVIDIA driver $DRIVER (570+ required for cuda-checkpoint)"
ERRORS=$((ERRORS + 1))
fi
else
fail "NVIDIA driver not found"
ERRORS=$((ERRORS + 1))
fi
# --- Podman ---
if command -v podman &>/dev/null; then
ok "podman $(podman --version | awk '{print $NF}')"
elif [ "$CHECK" = "true" ]; then
fail "podman not installed"
ERRORS=$((ERRORS + 1))
else
warn "Installing podman..."
sudo apt-get update -qq && sudo apt-get install -y -qq podman
ok "podman installed"
fi
# --- CRIU ---
if command -v criu &>/dev/null; then
ok "criu $(criu --version | head -1 | awk '{print $NF}')"
elif [ "$CHECK" = "true" ]; then
fail "criu not installed"
ERRORS=$((ERRORS + 1))
else
warn "Installing CRIU from PPA..."
sudo add-apt-repository -y ppa:criu/ppa 2>/dev/null
sudo apt-get update -qq && sudo apt-get install -y -qq criu
ok "criu installed"
fi
# --- CRIU default config (link-remap) ---
if [ -f /etc/criu/default.conf ] && grep -q "link-remap" /etc/criu/default.conf; then
ok "CRIU link-remap enabled"
elif [ "$CHECK" = "true" ]; then
fail "CRIU link-remap not configured (/etc/criu/default.conf)"
ERRORS=$((ERRORS + 1))
else
warn "Configuring CRIU link-remap..."
sudo mkdir -p /etc/criu
echo "link-remap" | sudo tee /etc/criu/default.conf >/dev/null
ok "CRIU link-remap configured"
fi
# --- nvidia-container-toolkit ---
if dpkg -l nvidia-container-toolkit &>/dev/null; then
ok "nvidia-container-toolkit installed"
elif [ "$CHECK" = "true" ]; then
fail "nvidia-container-toolkit not installed"
ERRORS=$((ERRORS + 1))
else
warn "Installing nvidia-container-toolkit..."
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null
sudo apt-get update -qq && sudo apt-get install -y -qq nvidia-container-toolkit
ok "nvidia-container-toolkit installed"
fi
# --- CDI spec ---
if [ -f /etc/cdi/nvidia.yaml ]; then
ok "NVIDIA CDI spec configured"
elif [ "$CHECK" = "true" ]; then
fail "NVIDIA CDI spec not generated (/etc/cdi/nvidia.yaml)"
ERRORS=$((ERRORS + 1))
else
warn "Generating NVIDIA CDI spec..."
sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml 2>/dev/null
ok "NVIDIA CDI spec generated"
fi
# --- cuda-checkpoint ---
if command -v cuda-checkpoint &>/dev/null; then
ok "cuda-checkpoint found"
else
fail "cuda-checkpoint not found in PATH (ships with NVIDIA driver 570+)"
ERRORS=$((ERRORS + 1))
fi
# --- Podman GPU test ---
if [ "$CHECK" = "true" ] || command -v podman &>/dev/null; then
if sudo podman run --rm --privileged --device nvidia.com/gpu=0 docker.io/nvidia/cuda:12.4.0-base-ubuntu22.04 nvidia-smi &>/dev/null; then
ok "podman GPU access works"
else
fail "podman GPU access failed"
ERRORS=$((ERRORS + 1))
fi
fi
echo ""
if [ "$ERRORS" -gt 0 ]; then
fail "$ERRORS issue(s) found"
exit 1
else
ok "All checks passed"
fi