Google Colab proをVSCodeでssh接続してV100(GPU)を早速回してみた。

はじめに

こんにちはNever$と申します。製造業でデータ分析系の業務についております(Kaggleも一応嗜んでおります)
Twitterで"VSCodeでColab-pro繋いだ~"と書いたら意外と反響がありましたのでやり方を書いてみたいと思います。参考になったと思った方は"LGTM"ボタンをポチッとお願いします。

※今回はあくまで繋いで回すまでを実験的にやってみただけですので、この記事を読んで何か不具合が発生しても自己責任でお願いします。
「ngrokつなげたら不正アクセスされた」系がもしかしたらあるかもしれませんのでご注意ください(この辺りの知識が疎いため憶測です…)

参考サイト

・Google Colab をSSH と VS Code で使う
丁寧に説明されていてわかりやすかったです。

準備するもの

・Colab pro : 契約願います。
・VSCode : ローカルにインストール願います。
・ngrok : これが今回一番重要なもので、colabに接続するのに必要です。無料枠でOKですので登録願います。ググると詳しくわかります。

作業開始〜colab上での作業

まずはColab上で適当にNotebookを作成してください。
イメージはそこにアクセスする感じになります。

1 : Colab上でGPUをON, メモリをハイメモリにする

2 : google driveをマウントする

colabからgoogle driveを読みこむために必要なコードです。

from google.colab import drive
drive.mount('/content/drive')
print('Mount OK')

3 : SSH接続のためのコードを入力

この辺りは参考サイト記載内容を丸写ししております。詳細については参考先を参照願います。
1セルごとに実行してください。

import random, string, urllib.request, json, getpass

パスワードの生成

password = ''.join(random.choice(string.ascii_letters + string.digits) for i in range(20))
password

ngrokのDL

! wget -q -c -nc https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
! unzip -qq -n ngrok-stable-linux-amd64.zip

SSHの設定

! apt-get install -qq -o=Dpkg::Use-Pty=0 openssh-server pwgen > /dev/null

SSH周りの設定

#Set root password
! echo root:$password | chpasswd
! mkdir -p /var/run/sshd
! echo "PermitRootLogin yes" >> /etc/ssh/sshd_config
! echo "PasswordAuthentication yes" >> /etc/ssh/sshd_config
! echo "LD_LIBRARY_PATH=/usr/lib64-nvidia" >> /root/.bashrc
! echo "export LD_LIBRARY_PATH" >> /root/.bashrc

get_ipython().system_raw('/usr/sbin/sshd -D &')

ngrokのAPIトークンを入力

authtoken="API Tokenを入力"

get_ipython().system_raw('./ngrok authtoken $authtoken && ./ngrok tcp 22 &')

これで一旦終わり

with urllib.request.urlopen('http://localhost:4040/api/tunnels') as response:
  data = json.loads(response.read().decode())
  (host, port) = data['tunnels'][0]['public_url'][6:].split(':')
  print(f'SSH command: ssh -p{port} root@{host}')

print(f'Root password: {password}')

そうすると以下の文字が出力されます。
これは次に使いますので残しといてください。

SSH command: ssh -pXXXXX root@YYYYYYYY
Root password: ZZZZZZ

VSCodeの設定

はじめにVSCode上の拡張機能である"Remote-SSH"を導入します。
下図赤枠のところをポチポチしてインストールしてください。

Remote-SSHの設定

colab接続用のSSH TARGETを設定

作成したSSH TARGETを選択し歯車をクリックしてconfingをいじります。
下記のような感じで書ます。YYYYYYYYとXXXXXは前述のものです。

Host google_colab_ssh
    HostName YYYYYYY
    User root
    Port XXXXX

SSH接続

いよいよ本番です。
該当のSSH TARGETをアクティブにして右のフォルダマーク?をクリック
上の方に選択肢が出てきますのでContinueを押して、前述のZZZZZZを入力。

そうすると入れます。
入った後は拡張機能からpythonでも入れて、Jupyter形式で開きGPUを確認してみましょう!

! nvidia-smi

V100握ってる!!!

せっかくなので昔のコンペで使用したコードを回してみる。

左のOpen Folderからdriveを開きます。
Pathは"/content/drive/MyDrive/~"以下個人の環境になります。

ちゃんと自分のコード上でもV100握ってる!!!(V100やたらと出るな笑)

テーブルデータで学習開始!(MoAコンペのResNetいじったやつです)
体感ではめっちゃ早くなってる気がします。

終わりに

VSCodeとColab proの連携は以上となります。
駆け足で作成したため、多少説明が不足している点があると思いますがご了承ください。
私自身もまだ理解していない点が多数ありますので引き続き環境構築に注力したいと思います。

最後まで読んでいただきありがとうございました。