手で積み上げた ML の venv は、壊れたときに同じ物を作り直せるかが分かりません。私たちも、ROCm の技術検証で画像や音楽などのモデルを動かしていた venv 2本について、pip freeze --all の控えは取っていましたが、控えから一括で作り直したことは一度もありませんでした。動いている venv に触らずに作り直し、確かめてから切り替えた手順をまとめます。

結論

  • 控えは2つです。pip freeze --all の出力と、torch 一式4本をハッシュ付きで書いた要件ファイルです。
  • 新しい venv は、今の venv に触らず、控えの指紋を名前に入れた別の場所に作ります。作ってから名前は変えません(venv の中のスクリプトの1行目に絶対パスが入るため)。
  • 作ったら pip freeze --all を控えと diff し、1行でも違えば使いません。さらに生成の種類ごと(画像・効果音・音楽・検査・3D・文章)に1回ずつ通し、同じ seed の画像が画素の差 0 で一致することを確かめました。
  • 切り替えは symlink の張り替えで行い、戻すときは前の実体へ向け直します。

試した環境

  • GPU: AMD Radeon RX 7900 XT(gfx1100、VRAM 20 GiB)
  • OS と ROCm: Ubuntu、ROCm 7.2.4
  • venv は2本。画像・効果音・VL などの venv と、音楽(ACE-Step 1.5)の venv です。transformers の版が噛み合わないので分けました
  • torch 2.13.0+rocm7.2 ほか4本、Python 3.12、pip 26.2.1
  • 生成は、検証の中で組んだ、ジョブを並べて GPU で1本ずつ回す仕組みを通して、この venv で走らせました

流れ

次の図は、控えから作って切り替えるまでの流れです。

控えから作って切り替えるまで

今の venv には最後の張り替えまで触らないので、途中で失敗しても、動いている物は壊れません。

控えを作る

1つ目は pip freeze --all の出力です。2つ目は、torch・torchvision・torchaudio・triton-rocm の4本をハッシュ付きで書いた要件ファイルです(ハッシュは私たちの環境の wheel の値です)。

--index-url https://download.pytorch.org/whl/rocm7.2
torch==2.13.0+rocm7.2 --hash=sha256:2ad8066cb9d3c00b2dc77ec5af998eed495e7240e724e8c28124d4d5faba52f0
torchvision==0.28.0+rocm7.2 --hash=sha256:5f54d604f0394075d92f61127ca0e062258b7a5196c267a5a7b67dbee7f034e6
torchaudio==2.11.0+rocm7.2 --hash=sha256:e0316d938346cfa4ff04b91405a91ecc50bd18781023eca5763b79f0f7b9bc8c
triton-rocm==3.7.1 --hash=sha256:b200ade2418b8450e4e3e69c33c804e2cc07d1e36dc9678ac8ac1dd65f9e10e1

ここで1つつまずきました。download.pytorch.org/whl/rocm7.2 の index は、torch 2.13.0+rocm7.2 の wheel にだけ(8本とも)sha256 を載せていませんでした。同じページでも、ほかの torch の版には載っています。そこで wheel 全体(6,227,849,000 B)を取って自分で計算し、Content-Length と配信元(R2)の ETag も一致することを確かめました。torchvision・torchaudio・triton-rocm は、index の sha256 と一致しました。

版を固定しないと、執筆時点では torch 2.14.1 + torchvision 0.29.1 が入り、今の環境とは別の組み合わせになります。

作る

作るスクリプトの要所です(場所は例に置き換えています)。

#!/bin/sh
# 控え(freeze とハッシュ付きの torch)から、venv を新しい場所に作る
#   使い方: sh <このスクリプト> <venv の名前>
set -eu
NAME=${1:?venv の名前}
L=$(cd "$(dirname "$0")" && pwd)            # 控えを置いたディレクトリ
MODELS=${MODELS:?モデル用のディスクの場所}
MOSS_SRC=${MOSS_SRC:?MOSS-TTS のソースの場所}
FREEZE=$L/$NAME.freeze.txt
SUM=$(cat "$L/torch-hashes.txt" "$FREEZE" | sha256sum | cut -c1-12)
DEST=$MODELS/venvs/$NAME-$SUM
[ ! -e "$DEST" ] || { echo "もうある: ${DEST}" >&2; exit 1; }

export PIP_CACHE_DIR=$MODELS/.cache/pip     # 既定の ~/.cache はルートのディスク
export PIP_DISABLE_PIP_VERSION_CHECK=1
mkdir -p "$MODELS/venvs" "$PIP_CACHE_DIR"
python3.12 -m venv "$DEST"                  # 作ってから名前を変えない
P=$DEST/bin/pip
"$P" install --quiet pip==26.2.1            # 24.0 は切れたダウンロードを再開できない
"$P" install --quiet -r "$L/torch-hashes.txt" --require-hashes --no-deps
# 残りは freeze から(torch の4本・-e の行・コメント行を除く)
grep -vE '^(torch|torchvision|torchaudio|triton-rocm)==|^-e |^#' "$FREEZE" > "$DEST.rest.txt"
"$P" install --quiet -r "$DEST.rest.txt"
if [ "$NAME" = venv ]; then                 # 効果音のコードは自分の pin が多いので --no-deps
  "$P" install --quiet --no-deps -e "$MOSS_SRC/moss_soundeffect_v2"
fi
rm -f "$DEST.rest.txt"

"$P" freeze --all > "$DEST.freeze.txt"
if ! diff "$FREEZE" "$DEST.freeze.txt"; then
  echo "控えと違う(上の差分)。${DEST} は使わない" >&2
  exit 1
fi
rm -f "$DEST.freeze.txt"
echo "作った: ${DEST}(控えと一致)"
  • 名前の指紋: 要件ファイルと freeze をつないだ sha256 の先頭 12 桁です。控えが変われば名前も変わり、同じ名前が既にあれば止まります。
  • pip を先に上げる: venv の既定の 24.0 は、途中で切れたダウンロードを再開できません。torch は 6.2 GB あります。
  • torch は --require-hashes --no-deps: --require-hashes は依存にもハッシュを求めるので、--no-deps が要ります。
  • 残りは freeze から: torch の4本・-e の行・コメント行を除いて入れます。diffusers の git の行はそのまま通り、pip の行で控えの版にそろいます。
  • 最後に diff: 1行でも違えば、その venv は使いません。

結果は、2本とも 16 GB ほどで、数分で作れました。freeze は控えと1行も違いませんでした。

確かめる

freeze が同じでも、動くかは別に確かめます。ジョブ管理を、新しいジョブを配らない状態にしたうえで、サービス用の利用者として、過去のジョブと同じ引数と環境で生成の種類ごとに1回ずつ通しました。生成を起こすスクリプトは環境変数で使う venv を選べるので、そこだけ新しい方に向けます。

  • 画像・効果音・音楽(turbo)・画像を見る VL と音の検査・3D(解像度 512)・文章(LLM)を1回ずつ
  • 全部 exit 0
  • 画像は、同じ seed の前の出力と画素の差 0

切り替える

ジョブを配らない状態のまま、symlink を張り替えます。今の venv が symlink でなく実ディレクトリのときは、最初の1回だけ実体を退避します。前の実体は消さずに残します。

cd <モデル用のディスク>
mv venv venvs/venv-old        # 最初の1回だけ(今の venv が実ディレクトリのとき)
ln -sfn venvs/venv-<指紋> venv.new && mv -T venv.new venv

張り替えたあと、ジョブの配りを戻して画像を1本作り、新しい venv で動くことを確かめました。

戻すときは、symlink を前の実体へ向け直します。前の実体(各 16 GB)は、新しい venv で問題なく動いているのを見てから消しました。消したあとに戻すときは、前の控えから同じ手順で作り直します。pip のキャッシュ(6.5 GB)は、作り直しに使うので残しました。

venv を直したいときも、今の venv に pip install せず、控えを直して新しく作り、確かめてから切り替えれば、動いている物を壊さずに済みます。

回転台の上で、新しい鉢植えが前に出て、古い鉢植えが後ろに残されている絵

確かめていないこと

  • 何か月か後にも、同じ版が PyPI と PyTorch の index から取れ続けるか。今回は取れました。pip のキャッシュを残したのはそのためです。
  • ほかの道具用の venv(変換やビルドの道具など)は、freeze の控えを置いただけで、控えから作り直したことはありません。

おわりに

控えは、作り直して差分 0 になるまで確かめて、初めて控えとして使えます。今の venv に触らずに別の場所に作り、出力の一致まで見てから symlink で切り替えれば、失敗しても動いている物は壊れません。

※ download.pytorch.org の index の載せ方と pip の振る舞いは執筆時点の物で、いまは変わっている可能性があります。