ssh 连接 Google Colab 服务器

Google Colab 提供了免费的服务器资源，可以使用 Google Colab 的 CPU，也可以使用 GPU 运行深度学习程序。然而，Google Colab 只提供了网页的操作界面，不能像平时操作服务器一样在真实的终端中输入命令。为此，本文提供一种借助中转服务器，使用 ssh 连接 Google Colab 服务器的方法。

实现思路

为了通过 ssh 连接 Google Colab 服务器，需要一个有公网 IP 的中转服务器。通过端口转发，将 Google Colab 服务器的 ssh 端口转发到中转服务器上，然后连接中转服务器，达到连接 Google Colab 服务器的目的。

实现步骤

将中转服务器的公钥加入 Google Colab 服务器

%%sh
mkdir -p ~/.ssh
echo '<PUBLIC_KEY>' >> ~/.ssh/authorized_keys

配置 Google Colab 服务器

%%sh
# Update system and install necessary packages
apt update > /dev/null
yes | unminimize > /dev/null
apt install -qq -o=Dpkg::Use-Pty=0 openssh-server pwgen net-tools psmisc pciutils > /dev/null

# Config ssh
ssh-keygen -t rsa -N "" -f ~/.ssh/id_rsa

# Config sshd
echo ListenAddress 127.0.0.1 >> /etc/ssh/sshd_config
mkdir -p /var/run/sshd
/usr/sbin/sshd

%%sh
# Install packages
apt install -qq -o=Dpkg::Use-Pty=0 zsh neofetch python-virtualenv nano python3.8 python3.8-dev python3.8-distutils python3.8-venv > /dev/null

# Set default shell to zsh
chsh -s `which zsh`

# Install oh-my-zsh
yes | sh -c "$(curl -fsSL https://raw.github.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"

# Customize zsh theme
wget -q -P /root/.oh-my-zsh/custom/themes https://gist.githubusercontent.com/ayaka14732/d15e5a583de03a3b310019f869cc2302/raw/118b7b46ffcf560bf85a35460d7df8343abc62c1/af-magic-%25E7%25B6%25BE.zsh-theme
sed -r -i 's/robbyrussell/af-magic-綾/' ~/.zshrc

# Install pip for Python 3.8
curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
python3.8 get-pip.py
rm -f get-pip.py

# Set timezone
echo 'export TZ=":Asia/Hong_Kong"' >> ~/.zshrc

# Set alias
echo 'alias python=python3.8
alias pip="python3.8 -m pip"' >> ~/.zshrc

将 Google Colab 服务器的公钥加入中转服务器

!cat ~/.ssh/id_rsa.pub

echo '<PUBLIC_KEY>' >> ~/.ssh/authorized_keys

将 Google Colab 服务器的 ssh 端口转发到中转服务器

%%sh --bg
ssh -fNT -R 7423:127.0.0.1:22 user@example.com

其中，user@example.com 要替换为中转服务器登录用户名和 IP 地址。

连接 Google Colab 服务器

ssh user@example.com

其中，user@example.com 要替换为中转服务器登录用户名和 IP 地址。

ssh -p7423 root@127.0.0.1

当重新启动 Google Colab 运行环境时，Google Colab 服务器的 host key 会发生变化，导致无法再次在中转服务器上登录 Google Colab 服务器。此时可以在中转服务器上运行以下命令，删除原来储存的 host key：

ssh-keygen -R "[127.0.0.1]:7423"

备考

查看 GPU 型号

import torch
print(torch.cuda.get_device_name())

如果 GPU 型号不理想，可以尝试重新启动 Google Colab 运行环境。

查看系统信息

查看 GPU 使用情况

LD_LIBRARY_PATH=/usr/lib64-nvidia watch -n 1 nvidia-smi

不能直接使用 export LD_LIBRARY_PATH=/usr/lib64-nvidia，这样会导致某些程序不能使用 GPU。

保持 ssh 连接

为了防止 ssh 连接一段时间后就中断，可以在 ssh 选项中加入：

ServerAliveInterval=60

挂载 Google Drive

在启动 8 小时左右，服务器资源会被回收。为了防止数据丢失，可以定期将中间结果保存到 Google Drive 中。将 Google Drive 挂载到 Google Colab 的方法是在 Google Colab 单元格中运行：

from google.colab import drive
drive.mount('/content/gdrive')

防止没有操作网页导致资源被回收

如果一段时间没有操作网页，服务器资源将被回收。为此，可以在网页上运行一段输出日志的代码，告诉 Google Colab 这个服务器还在用：

import time
import datetime

with open('out.log', 'w') as f:
    for i in range(10000):
        print(datetime.datetime.now(), 'Doing hard work here i =', i)
        print(datetime.datetime.now(), 'Doing hard work here i =', i, file=f)
        time.sleep(1000)

简介