如何在Google Compute Engine上运行TensorFlow GPU容器？

Question

我想在拥有GPU加速器的Google计算引擎上运行TensorFlow容器。

尝试了以下命令：

gcloud compute instances create-with-container job-name \
  --machine-type=n1-standard-4 \
  --accelerator=type=nvidia-tesla-k80 \
  --image-project=deeplearning-platform-release \
  --image-family=common-container \
  --container image gcr/io/my-container \
  --container-arg="--container-arguments=xxxx"

但是收到了警告

警告: 此容器部署机制需要一个Container-Optimized OS镜像才能运行。从cos-cloud项目（cost-stable、cos-beta、cos-dev镜像系列）中选择一个镜像。

我还尝试了来自cos-cloud项目的系统镜像，但似乎没有CUDA驱动程序，因为tensorflow记录了警告cuInit failed。

想知道在具有GPU支持的Google计算引擎上运行tensorflow容器的正确方法是什么？

- liuyl

Answer 1

你可以在 deeplearningvm 的 startup-script 中使用 docker run 命令来运行你的容器。

gcloud beta compute instances create deeplearningvm-$(date +"%Y%m%d-%H%M%S") \
--zone=us-central1-c \
--machine-type=n1-standard-8 \
--subnet=default \
--service-account=<your google service account> \
--scopes='https://www.googleapis.com/auth/cloud-platform' \
--accelerator=type=nvidia-tesla-k80,count=1 \
--image-project=deeplearning-platform-release \
--image-family=tf-latest-gpu \
--maintenance-policy=TERMINATE \
--metadata=install-nvidia-driver=True,startup-script='#!/bin/bash

# Check the driver until installed
while ! [[ -x "$(command -v nvidia-smi)" ]];
do
  echo "sleep to check"
  sleep 5s
done
echo "nvidia-smi is installed"

gcloud auth configure-docker
echo "Docker run with GPUs"
docker run --gpus all --log-driver=gcplogs --rm gcr.io/<your container>

echo "Kill VM $(hostname)"
gcloud compute instances delete $(hostname) --zone \
$(curl -H Metadata-Flavor:Google http://metadata.google.internal/computeMetadata/v1/instance/zone -s | cut -d/ -f4) -q

'

由于安装NVIDIA驱动程序需要几分钟的时间，请在安装完成后再启动容器。Compute Engine会在第一次启动时加载最新的稳定驱动程序，并执行必要的步骤（包括最后一次重启以激活驱动程序）。在此期间，您将无法通过SSH登录到您的机器。安装完成后，为确保驱动程序安装成功，您可以通过SSH登录并运行nvidia-smi命令。详情请参见：https://cloud.google.com/ai-platform/deep-learning-vm/docs/tensorflow_start_instance#creating_a_tensorflow_instance_from_the_command_line。

- northtree

---

能否在深度学习平台中使用自定义的Docker镜像？ - Max Voitko

@MaxV 当然可以。请参考我的脚本中的 docker run 命令。 - northtree

Answer 2

你考虑过在GKE上使用Cloud TPU吗？

这个页面描述了如何设置带有GPU的GKE集群。

- Ernesto U

---

@liuyl，你搞清楚怎么在没有GKE的情况下运行了吗？ - northtree

很抱歉，我已经转向使用GKE了。 - liuyl

@liuyl，感谢您的更新。我想知道是否可以通过GKE轻松设置。 - northtree