Google Colaboratoryにsshログインをしてお手軽GPU実験環境を作ってみた

Google Colaboratoryの制限事項の改訂で現在はSSH接続による利用は禁止されています。

Googleが無料で公開しているGPU、TPU上でJupyter Notebookライクなノートブックを実行出来るGoogle Colaboratoryは無料とは思えない程機能が充実しています。あるセミナーでsshでログインをして利用出来ることを知りました。最近CUDA Cに関心が有り、無料・簡単・早いの三拍子揃った勉強用環境は作れないかと思っていたので、以下の方法なら三拍子揃った環境が構築出来るのではないかと思いやってみました。

GPUマシン自作やクラウドサービスでは駄目か

GPUを利用出来る環境の構築にはGPUを搭載したPCの自作やクラウドサービス(AWSやGCP）をレンタルする方法が有りますが、気軽に試すには初期投資が掛かり過ぎたり、停止を忘れると高額な請求がカードにくる危険性が有ったりと勉強内容とは関係の無いところで頭痛の種が出来てしまうのである程度勉強が進んでから導入を検討することにします。

前準備 : Googleアカウントの登録

Google Colaboratoryの利用にはGoogleアカウントが必要です。事前にアカウントを取得しておきましょう。

Google Colaboratoryにログイン

Google Colaboratoryにログインします。

ngrokアカウントの登録

Google Colaboratoryをsshログインするためにngrokというサービスを利用します。サービス内容等の詳細は以下を参考にしてください。アカウント登録後はログインをしておきます。
https://ngrok.com/
ngrokを使ってお手軽に開発環境のWebサーバを外から接続できるようにしよう

GPUノートブックの作成

以下のページが参考になります。
【秒速で無料GPUを使う】TensorFow(Keras)/PyTorch/Chainer環境構築 on Colaboratory

sshログインの為の前準備

認証情報取得用Pythonシェルを実行

セル(コマンド欄)に以下の内容を記載し、セルを実行します。

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

#Download ngrok
! wget -q -c -nc https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
! unzip -qq -n ngrok-stable-linux-amd64.zip
#Setup sshd
! apt-get install -qq -o=Dpkg::Use-Pty=0 openssh-server pwgen > /dev/null
#Set root password
! echo root:$password | chpasswd
! mkdir -p /var/run/sshd
! echo "PermitRootLogin yes" >> /etc/ssh/sshd_config
! echo "PubkeyAuthentication 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

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

#Ask token
print("Copy authtoken from https://dashboard.ngrok.com/auth")
import getpass
authtoken = getpass.getpass()

#Create tunnel
get_ipython().system_raw('./ngrok authtoken $authtoken && ./ngrok tcp 22 &')
#Print root password
print("Root password: {}".format(password))
#Get public address
! curl -s http://localhost:4040/api/tunnels | python3 -c \
    "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"

実行後暫くすると以下の様な内容が表示されます。URLをクリックします。

パスワードが表示されます。Copyボタンをクリックします。

ノートブックにパスワードをペースト

上の行程でコピーしたパスワードを先程のノートブックのセルのURL下に表示されているテキストボックスにペーストし、Enterキーを押します。

暫くすると認証情報が表示されます。

sshpassのインストール

セミナーで聞いた方法や参考にしたWebページでは一貫して共通鍵認証でsshログインを実現する方法が紹介されていました。しかし、Google Colaboratoryは連続実行時間に制限が有り、12時間毎にセッションがリセットされます。リセット後は、毎回キーファイルを登録し直さなければなりません。故に今回は、パスワード認証でログインすることにしました。でも、毎回のパスワード入力は面倒ですのでsshpassをインストールしてパスワード入力を自動化します。

Ubuntuの場合

以下のコマンドを実行します。

sudo apt install -y sshpass

macOSの場合(Homebrewを利用)

Homebrewでインストールするには以下のコマンドを実行します。

brew install http://git.io/sshpass.rb

参考

• macOS に Homebrew で sshpass をインストールする
  近年のMacはHomebrewでのインストールが非推奨になっています。以下のの記事を参考にして下さい。
• Homebrewでインストールが非推奨!?になったsshpassはソースコードビルドでインストールしよう

WindowsはWindows10からsshに対応しましたが、sshpassは非対応の様です。ただ、類似の機能が有る様です。また、WSL(Windows Subsystem for Linux)を利用すればUbuntuでのインストール方法が使えます。

ログイン用シェルの用意

ターミナルを開いて適当なシェルファイルを作成し、以下の内容を記載します。chmod +x hogehoge.shを実行しておけば、ターミナル上で./hogehoge.shで実行出来ます。

#!/bin/bash

port=
passwd=
sshpass -p $passwd ssh -o "StrictHostKeyChecking no" root@tcp://0.tcp.ngrok.io -p $port

このシェルファイルに先程の認証情報から必要な情報を入力します。認証情報の画像を再掲します。

今回の例では、portには11986が入ります。port=11986、passwdにはpassword:項の右側の値(塗りつぶしている箇所)を入力します。これで完成です。ポートとパスワードについてはセッションがリセットされる毎に上述のPythonのスクリプトをノートブック上で実行して認証情報を取得、値の変更が必要になります。故に引数で渡せる様にしても良いかと思われます。

sshログイン

完成したシェルファイルを実行して実際にログインしてみましょう。

ログインに成功すると以下の様な画面が表示されます

Warning: Permanently added '[tcp://0.tcp.ngrok.io]:15227,[3.19.114.185]:15227' (ECDSA) to the list of known hosts.
Welcome to Ubuntu 18.04.3 LTS (GNU/Linux 4.14.137+ x86_64)

 * Documentation:  https://help.ubuntu.com
 * Management:     https://landscape.canonical.com
 * Support:        https://ubuntu.com/advantage
This system has been minimized by removing packages and content that are
not required on a system that users do not log into.

To restore this content, you can run the 'unminimize' command.

The programs included with the Ubuntu system are free software;
the exact distribution terms for each program are described in the
individual files in /usr/share/doc/*/copyright.

Ubuntu comes with ABSOLUTELY NO WARRANTY, to the extent permitted by
applicable law.

root@54eb64b85a06:~# 

CUDA実行環境のPath登録

この環境にはCUDA用の環境がインストールされていますが、PATH登録が出来ていません。故に、nvcc -vと叩いても以下の様な内容が出力されます。

root@54eb64b85a06:~# nvcc -V
-bash: nvcc: command not found

CUDA C実行環境PATH登録用スクリプトの作成

以下の内容を記載したシェルファイルを作成します。再利用することを前提に作成しているので、GitHub等にリポジトリを作成して登録しておき、利用時には毎回cloneしてくる運用をすると良いかと思われます。

#!/bin/bash

echo "export PATH="/usr/local/cuda/bin:$PATH"" >> ~/.bashrc
echo "export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"" >> ~/.bashrc

ログインシェルと同様に実行します。そして、. .bashrc(またはsource .bashrc)を実行し、.bashrcの再読込を行います。

root@54eb64b85a06:~# nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on Sat_Aug_25_21:08:01_CDT_2018
Cuda compilation tools, release 10.0, V10.0.130

正しくPATHが通り、nvccコンパイラが実行出来る様になりました。これで環境は完成です。連続実行時間は12時間が上限となっている様ですが、個人のお試し利用なら十分実用に耐えうるでしょう。

まとめ

Google Colaboratoryをsshログインしてお手軽GPU実験環境を構築しました。今後実際にCUDA Cを用いたプログラムのコンパイル、実行についての内容も追加していこうと思います。

Reference

• TadaoYamaokaの日記
• sshでGoogle Colaboratoryを使う
• 【秒速で無料GPUを使う】深層学習実践Tips on Colaboratory