Google Colab ProのGPUガチャを快適にまわす

背景

最近Colab Pro +に課金をしたのですが，最も強いA100は人気が高いせいか時々接続ができずに下位のGPUに接続されるということがあります。(通称GPUガチャ)
その場合，セッションをリセットしリソースを変更する必要があるのですが，複数回これを実行するとなると結構めんどくさいです。
L4も悪くはない気がしているのですがせっかく課金したなら強いGPUを使いたいですよね。そのためセルを実行するだけでイイ感じにリセットしてくれるコードを書きました。

やりかた

下記コードをセルの最上部に張り付けて実行するのみです。

import subprocess
import os
import re

def get_gpu_memory():
    result = subprocess.run(
        ['nvidia-smi', '--query-gpu=memory.total', '--format=csv,nounits,noheader'],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        print(f"Error querying GPU memory: {result.stderr}")
        return None
    return int(result.stdout.strip().split('\n')[0])

vram = get_gpu_memory()

if vram is not None:
    print(f"Detected GPU VRAM: {vram} MB")
    if vram < 40960:  # 40GB = 40960MB
        print("VRAM is less than 40GB. Restarting the session...")
        os.kill(os.getpid(), 9)
    else:
        print("VRAM is 40GB or more. No need to restart the session.")
else:
    print("Unable to detect GPU VRAM.")

やってること自体は単純で，GPUのメモリ量を取得して40GB未満なリセットしてます。24年6月時点ではA100以外がVRAM40GB未満なのでA100に接続できれば以降のセルも正しく実行できます。

おわりに

A100以外を使いたい場合，このコードを実行するとセッションリセットされるので消すかコメントアウトを忘れずに。快適にGPUガチャを回しましょう。