Google Colabのランタイムをコードで終了させる方法

はじめに

Colab pro, Colab Pro+などのプランが変更になり、利用時間でコンピューティングユニットを消費するようになりました。以前は気にしていませんでしたが、従量制になると「処理が終了したら切断したい」と考えてしまいます。

ここでは、Google Colabのランタイムをコードから終了させる方法を紹介。

やり方

最後に終了させたい場合

色々な方が書かれていますが、非常にシンプルです

from google.colab import runtime

runtime.unassign()

これを最後に追加するだけ。

エラー時に終了させたい

こちらもやりたいですが、簡単な方法が見つからなかったです。とりあえず例外処理を使って、こんな感じで書けば一応、終了します。

from google.colab import runtime
import sys

try:
  x = np.array([10,20,30])
except Exception as e:
  print(e.args)
  print(sys.exc_info())
  runtime.unassign()

("name 'np' is not defined",)
(<class 'NameError'>, NameError("name 'np' is not defined"), <traceback object at 0x7f155008f980>)

この例では`import numpy as np`が無いのでエラーになりますが、その場合はエラーを表示して終了させることができます。ただ、毎ブロックこれを書くのは面倒。何かいい方法があれば教えて欲しいです。

もっと手軽な方法は、次の方法になります。

エラー時に終了させるためのコード（2024.4.16追記）

コメントにあるように、エラー時に終了させる方法を教えていただき、動作の確認ができましたの追記しておきます。

以下のコードを最初に貼り付けておけば、エラー時に終了させることが可能です。

wolさん、ありがとうございます。

# prompt: エラー時にGoogle Colabを強制終了する
from IPython.core.ultratb import AutoFormattedTB
from google.colab import runtime
import time
import sys

# IPythonのデフォルトのエラー表示機能を初期化
itb = AutoFormattedTB(mode="Verbose", color_scheme='Neutral', tb_offset=1)

def custom_exc(shell, etype, evalue, tb, tb_offset=None):
    # エラーメッセージを出力
    print("エラーが発生しため処理が中断されました。")
    print(etype, evalue, tb)

　　 #
    # 必要ならここに結果の保存処理などを入れておいても良いかも
    # 

    time.sleep(3)

    # Google Colabを強制終了
    runtime.unassign()
    time.sleep(1)
    runtime.force_kill_kernel()

# カスタム例外ハンドラをIPythonに設定
get_ipython().set_custom_exc((Exception,), custom_exc)

上記のコードを入れて意図的にエラーを発生させたとろこうまく切断できました。

x = np.array([10,20,30])

エラーが発生しました。
<class 'NameError'> name 'a' is not defined <traceback object at 0x7aa516134b00>