GoogleCloud x Terraformで生成AI実験用サーバーを構築する

最近は生成AIをつかっていろいろ遊んで実験しています。株式会社ウェイブの渡邉です。

生成AI系のツールを少し試すにはGoogle Colabは簡単でよいですが、いろいろデータを入れたり、毎日動かすにはどうしても不自由です。

本記事では、Google Cloud上に生成AIを動かすためのGPUサーバーをTerraformを使って構築した際のポイントをメモしておきます。

GPUサーバーの利用申請をする

Terraformの設定が済んで、GPUサーバーを立ち上げようとすると、制限に引っかかってるよと警告が...
GPUサーバーを利用したい場合は、別途申請が必要でした。
複数の項目で申請が必要なのでご注意ください。
例えばT4を使いたい場合に「GPUs(all regions)」、「Commited NVIDIA T4 GPUs」のそれぞれで申請が必要です。

「各リージョンで作成するGPUモデルに対する割り当て」に加え、「すべてのゾーンにおけるすべてのタイプのGPUの総数に対するグローバル割り当て」もリクエストする必要
GCPのDeep Learning VMインスタンスを立ち上げる #DeepLearning - Qiita

申請の説明のところは「機械学習の開発のため」と書いて申請しました。

Terraformのファイル

下記の設定をしています。

• 特定のIPからのアクセスのみ受け付ける
• SSHログインできる
• インスタンスタイプを柔軟に変更できる
• 指定時間に停止される

provider "google" {
  project = "xxxx"
  region  = "us-central1"
  zone    = "us-central1-c"
}

# 各種パラメータを設定する
locals {
  vm_instance = {
    t4      = { machine_type = "n1-standard-8", accelerator_type = "nvidia-tesla-t4" },
    l4      = { machine_type = "g2-standard-8", accelerator_type = "nvidia-l4" },
    a100    = { machine_type = "a2-highgpu-1g", accelerator_type = "nvidia-tesla-a100" },
    a100-80 = { machine_type = "a2-ultragpu-1g", accelerator_type = "nvidia-a100-80gb" }
  }
  env_name         = "dev"
  vm_instance_name = "a100-80"
  user_name        = "fugafuga"
  public_key       = "ssh-rsa xxxx"
  vm_disk_size     = 300
}

# VPCネットワークの作成
resource "google_compute_network" "vpc_network" {
  name                    = "${locale.env_name}-network"
  auto_create_subnetworks = false
}

# サブネットの作成
resource "google_compute_subnetwork" "subnet" {
  name          = "${locale.env_name}-subnet"
  ip_cidr_range = "10.0.1.0/24"
  region        = "us-central1"
  network       = google_compute_network.vpc_network.id
}

# ファイアウォールルールの作成
# 許可するIPのみにアクセスを制限する
resource "google_compute_firewall" "restrict_ip" {
  name    = "${locale.env_name}-allow-ssh-from-specific-ip"
  network = google_compute_network.vpc_network.name

  allow {
    protocol = "tcp"
    ports    = ["22", "80", "443"]
  }

  # 許可するIPアドレス範囲を指定
  source_ranges = ["xxx.xxx.xxx.xxx/32"]
}

# 固定IPの生成
resource "google_compute_address" "static" {
  name = "${locale.env_name}-ipv4-address"
}

# 自動停止の設定
resource "google_compute_resource_policy" "stop_schedule" {
  name   = "stop-schedule-${locale.env_name}"
  region = "us-central1"

  instance_schedule_policy {
    vm_stop_schedule {
      schedule = "0 20 * * *" # 毎日20時(JST)に停止
    }
    time_zone = "Asia/Tokyo"
  }
}

# インスタンスのブートディスクを別途定義。これにより、インスタンスを再作成してもデータが保持される。
resource "google_compute_disk" "vm_disk" {
  name  = locale.env_name
  image = "https://www.googleapis.com/compute/v1/projects/deeplearning-platform-release/global/images/common-cu121-v20240514-debian-11-py310"
  size  = locale.vm_disk_size
}

# https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_instance
resource "google_compute_instance" "vm_instance" {
  name = locale.env_name

  machine_type = local.vm_instance[locale.vm_instance_name].machine_type
  metadata = {
    "block-project-ssh-keys" = "true"
    "ssh-keys"               = "${locale.user_name}:${locale.public_key}\n"
  }

  guest_accelerator {
    type  = local.vm_instance[locale.vm_instance_name].accelerator_type
    count = 1
  }

  boot_disk {
    auto_delete = false
    source      = google_compute_disk.vm_disk.self_link
  }

  network_interface {
    subnetwork = google_compute_subnetwork.subnet.id
    access_config {
      nat_ip = google_compute_address.static.address
    }
  }

  resource_policies = [google_compute_resource_policy.stop_schedule.id]

  scheduling {
    on_host_maintenance = "TERMINATE"

    # SPOTインスタンスを使うことでコスト削減
    provisioning_model          = "SPOT"
    preemptible                 = true
    automatic_restart           = false
    instance_termination_action = "STOP"
  }

  # 必要なライブラリがあれば変更
  metadata_startup_script = <<-EOF
    #!/bin/bash
    apt-get update -y
    apt install google-perftools git-lfs locales locales-all
  EOF
}

まとめ

以上のTerraformのコードを使用することで、GoogleCloud上にGPUを搭載した生成AI実験用サーバーを簡単に構築できました。必要に応じてGPUタイプを変更するも可能です。