Google Colaboratoryのランタイムを自動で終了&Slackへの通知

Others

Google Colaboratoryを使って時間のかかる処理を行う際に、コードの実行が終了したタイミングで自動でColabのランタイムを終了し、それをSlackへ通知するようにしてみました。

概要

以下が今回作成したサンプルコードの全体になります。これをGoogle Colaboratory上で実行すると、main関数を実行しながら経過を通知、終了した際には自動的にランタイムを終了するようになっています。また、実行中にエラーが発生した場合には、エラー内容をSlackへ送信するようにしています。

通知を送る際には、send_msgという関数を作成し、そこへテキストを与えるとSlackへ通知が行くようにしています。Google Colaboratoryのランタイム終了には、google.colab.runtime.unassign()を使用しています。

import time
import sys
import traceback
from google.colab import runtime

import json
import requests

# 任意のメッセージを通知する関数
def send_msg(message):
    # メッセージを飛ばすURL
    url = 'https://hooks.slack.com/services/XXXXXXXXXXX/XXXXXXXXXX/XXXXXXXXXXXXXXXXXXXXXXX' 
    
    data = json.dumps({'text': message})
    headers = {'content-type': 'application/json'}
    requests.post(url, data=data, headers=headers)


def main():
  for epoch in range(3):
    time.sleep(1)
    value = str(epoch+1)
    send_msg(f'{value}エポック')

try:
  send_msg('Start')

  # メインの処理
  main()

  # Google Colaboratoryのランタイムを終了
  runtimee.unassign()
  send_msg(f'Google Colabのランタイム終了')
  
except:
  # エラー発生時に内容をSlackへ通知
  send_msg('----- ERROR! -----\n'+traceback.format_exc())

Google Colaboratoryのランタイムを終了

Google Colaboratory Pro/Pro+では、コンピューティングユニットの分だけGPUを使用することができます。ランタイムに接続している間は、このコンピューティングユニットが消費されてしまうようなので、コードの実行が終わったら自動的にランタイムを終了することで、コンピューティングユニットの消費を最小限に抑えることができそうです。

以下のツイートによると、2022年9月からgoogle.colab.runtime.unassign()により、Google Colaboratoryのランタイムが終了させることができるようになったそうです。今回はこちらを使用してみました。

Slackへ通知を送る

PythonからSlackへ通知を送るには、Incoming Webhookを用いたSlackアプリを作成し、そのアプリで設定したURLを使用しています。

その仕事、Slack で。
Slack のことを学べるようサポートします。役立つ記事やビデオのコレクションを参照して、質問の答えをすばやく見つけましょう。

Slack Appの作成

以下の手順で、Slackへ通知を送るSlackアプリを作成します。

  1. https://api.slack.com/appsCreate New App を押し、From scratch を選択。
  2. App Name へアプリ名を入力し、Select a workspace からワークスペースを指定。
  3. Incoming Webhook を有効化する。
  4. 画面下の Add New Webhook to Workspace を押し、通知を送りたいチャンネルを指定してアプリへ権限を許可。

以上の手順を実行すると、Webhook URLが作成されます。このURLは、通知を送るチャンネルごとに固有に作成されるURLとなっています。そのため、別のチャンネルへ送信したい場合には、別途URLを作成する必要があります。

Slackへ通知を送る

作成したWebhook URLを用いて、Slackへ通知を送信します。Pythonのコードは以下の記事を参考にさせていただきました。

import json
import requests

# 任意のメッセージを通知する関数
def send_msg(message):
    # 作成したWebhook URL
    url = 'https://hooks.slack.com/services/XXXXXXXXXXX/XXXXXXXXXX/XXXXXXXXXXXXXXXXXXXXXXX' 

    data = json.dumps({'text': message})
    headers = {'content-type': 'application/json'}
    requests.post(url, data=data, headers=headers)

send_msg('Test')

上記コードを実行すると、しっかりSlackへ通知が送られてきました。アプリのアイコンを設定するとそれっぽい感じになりますね。

さいごに

今回作成したSlackアプリを使って、自分はDNNの学習時にエポックごとにロスやaccuracy等を送信させたり、エラーでコードが停止したことを通知させようと思っています。

また、Google Colaboratoryのランタイムを自動で終了させることができたため、学習が終了したらSlackへ通知して、自動的にランタイムを終了させるという使い方をしようと思っています。

もしどなたかの参考になれば幸いです。

コメント

タイトルとURLをコピーしました